Merge pull request #333 from zhijxu-MS/push_branch

bug fixe fill_op, enhance tf.where, enhance gather shape inference

Author: nbcsm

tests/test_backend.py

@@ -1502,6 +1502,16 @@ def test_where_scalar(self):
         _ = tf.identity(picks, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @check_opset_min_version(9, "where")
+    def test_where_with_cond_only(self):
+        for np_type, tf_type in [(np.int32, tf.int32), (np.float32, tf.float32)]:
+            x_val = np.random.randint(0, 2, size=[10, 20, 30]).astype(np_type)
+            x = tf.placeholder(tf_type, shape=[None] * x_val.ndim, name=_TFINPUT)
+            res = tf.where(x)
+            _ = tf.identity(res, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val})
+            tf.reset_default_graph()
+
     @check_opset_min_version(6, "cast")
     def test_shape_int32(self):
         x_val = np.array([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]], dtype=np.float32)

tf2onnx/optimizer/transpose_optimizer.py

@@ -398,11 +398,11 @@ def _maxmin_handler(self, trans, node):
         for i in all_other_inputs:
             target_node = self._g.get_node_by_output(i)
             numpy_val = target_node.get_tensor_value(as_list=False)
-            rank = np.rank(numpy_val)
+            rank = numpy_val.ndim
             if rank == 4:
                 transposed_val = np.transpose(numpy_val, (0, 3, 1, 2))
                 target_node.set_tensor_value(transposed_val)
-            elif rank == 1:  #  scalar
+            elif rank == 1:  # scalar
                 # do nothing
                 pass
             else:

tf2onnx/shape_inference.py

@@ -141,6 +141,17 @@ def infer_shape_for_node(g, node):
         log.debug("set ConcatV2 node [%s] with new shape %s", node.output[0], new_shape)
         return True
 
+    if node.type == "Gather":
+        # uses the follwing link to know how to infer shape of output
+        # https://www.tensorflow.org/api_docs/python/tf/gather
+        shape_params = g.get_shape(node.input[0])
+        shape_indices = g.get_shape(node.input[1])
+        axis = node.input[2].get_tensor_value()
+
+        shape = shape_params[:axis] + shape_indices + shape_indices[axis+1:]
+        g.set_shape(node.output[0], shape)
+        return True
+
     if node.type in ["All", "Any", "Min"]:
         axis_node = node.inputs[1]
         axis = axis_node.get_tensor_value()

tf2onnx/tfonnx.py

@@ -1463,17 +1463,16 @@ def fill_op(ctx, node, name, args):
     # both shape and value in tensorflow are passed as tensor.
     # In onnx the value is an attribute so we need to fetch the value as const which
     # sooner or later will be a problem for tensorflow-onnx.
-    shape = ctx.get_shape(node.output[0])
-    utils.make_sure(all(i >= 0 for i in shape), "shape attr should not be less than zero")
-    value = node.inputs[1].get_tensor_value()
-    value_proto = numpy_helper.from_array(node.inputs[1].get_tensor_value(as_list=False))
-    dtype = value_proto.data_type
-    # onnx spec says value MUST be float.
-    node.set_attr("value", float(value))
-    node.set_attr("shape", shape)
-    node.set_attr("dtype", dtype)
-    del node.input[:]
-    return node
+    # ConstantOfShape in onnxruntime only support int64, so insert cast op
+    input_dtype_is_int64 = utils.ONNX_TO_NUMPY_DTYPE[ctx.get_dtype(node.input[0])] == np.int64
+    if not input_dtype_is_int64:
+        cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0], to=onnx_pb.TensorProto.INT64)
+    dtype = ctx.get_dtype(node.output[0])
+    value = np.array([node.inputs[1].get_tensor_value()]).astype(utils.ONNX_TO_NUMPY_DTYPE[dtype])
+    value_proto = numpy_helper.from_array(value)
+    node.set_attr("value", value_proto)
+    del node.input[1]
+    return [node] if input_dtype_is_int64 else [node, cast_node]
 
 
 def reverse_op8(ctx, node, name, args):
@@ -1723,6 +1722,18 @@ def logical_compare_op(ctx, node, name, args):
     return nodes
 
 
+def where_op(ctx, node, name, args):
+    # T_y output = Where(T_x condition), return indices of elements whose value are True
+    node.type = "NonZero"
+    # in onnx, indices are returned in this way [[ind_a_0, ind_b_0, ...], [ind_a_1, ind_b_1,...]];
+    # while in tf, the result will be [[ind_a_0, ind_a_1, ...], [ind_b_0, ind_b_1, ...], ...]
+    # this is the reason a transpose node inserted here.
+    transpose_node = ctx.insert_new_node_on_output("Transpose", node.output[0], name=utils.make_name("where_op_added"))
+    ctx.copy_shape(node.output[0], transpose_node.output[0])
+    ctx.copy_dtype(node.output[0], transpose_node.output[0])
+    return [node, transpose_node]
+
+
 # map tensorflow ops to onnx ops. The format below is
 # "TFOP": func_to_map, ["OnnxOp", ...]
 #
@@ -1894,6 +1905,7 @@ def logical_compare_op(ctx, node, name, args):
     "Less": (logical_compare_op, []),
     "ResizeBilinear": (upsample_op9, ["Upsample", "linear"]),
     "ResizeNearestNeighbor": (upsample_op9, ["Upsample", "nearest"]),
+    "Where": (where_op, []),
 }
 
 _OPSETS = [