pass only int64 as shape to tile, exapand dims op

Author: guschmue

tests/test_backend.py

@@ -173,12 +173,15 @@ def test_trig_ops(self):
 
     @unittest.skipIf(BACKEND in ["caffe2", "onnxmsrt"], "not supported correctly in caffe2")
     def test_multinomial(self):
-        x_val = make_xval([3, 4])
+        shape = [2, 10]
+        x_val = np.ones(np.prod(shape)).astype("float32").reshape(shape)
         x = tf.placeholder(tf.float32, shape=x_val.shape, name=_TFINPUT)
         op = tf.multinomial(x, 2, output_dtype=tf.int64)
         output = tf.identity(op, name=_TFOUTPUT)
         actual, expected = self._run(output, {x: x_val}, {_INPUT: x_val})
-        self.assertAllClose(expected, actual, rtol=1e-06)
+        # since returned indexes are random we can only check type and shape
+        self.assertEqual(expected.dtype, actual.dtype)
+        self.assertAllClose(expected.shape, actual.shape)
 
     def test_maxppol(self):
         x_val = make_xval((1, 4, 4, 1))

tests/test_graph.py

@@ -14,6 +14,9 @@
 from tf2onnx.graph_matcher import *
 
 
+_TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
+
+
 def onnx_to_graphviz(g):
     g2 = gv.Digraph()
     for node in g.get_nodes():
@@ -266,7 +269,7 @@ def print_handler(ctx, node, name, args):
             # becomes:
             #   T output = Identity(T Input)
             node.type = "Identity"
-            node.domain = "tf"
+            node.domain = _TENSORFLOW_DOMAIN
             del node.input[1:]
             return node
 
@@ -276,7 +279,7 @@ def print_handler(ctx, node, name, args):
             _ = tf.identity(x_, name="output")
             g = process_tf_graph(sess.graph,
                                  custom_op_handlers={"Print": print_handler},
-                                 extra_opset=helper.make_opsetid("tf", 1))
+                                 extra_opset=helper.make_opsetid(_TENSORFLOW_DOMAIN, 1))
             self.assertEqual(
                 'digraph { Print [op_type=Identity] output [op_type=Identity] input1:0 -> Print Print:0 -> output }',
                 onnx_to_graphviz(g))

tf2onnx/convert.py

@@ -15,6 +15,9 @@
 from onnx import helper
 
 
+_TENSORFLOW_DOMAIN = "ai.onnx.converters.tensorflow"
+
+
 def get_args():
     """Parse commandline."""
     parser = argparse.ArgumentParser()
@@ -45,7 +48,7 @@ def get_args():
 
 
 def default_custom_op_handler(ctx, node, name, args):
-    node.domain = "tf"
+    node.domain = _TENSORFLOW_DOMAIN
     return node
 
 
@@ -61,7 +64,7 @@ def main():
     if args.custom_ops:
         # default custom ops for tensorflow-onnx are in the "tf" namespace
         custom_ops = {op: default_custom_op_handler for op in args.custom_ops.split(",")}
-        extra_opset = [helper.make_opsetid("tf", 1)]
+        extra_opset = [helper.make_opsetid(_TENSORFLOW_DOMAIN, 1)]
     else:
         custom_ops = {}
         extra_opset = None

tf2onnx/tfonnx.py

@@ -105,8 +105,25 @@ def tensorflow_to_onnx(graph):
     return onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes
 
 
-# pylint: disable=W0613,C0111,W0612
+def _convert_shapenode_to_int64(ctx, node, input_number):
+    shape_node = node.inputs[1]
+    name = node.input[1]
+    if shape_node.is_const():
+        # if it is a const, change the const to be int64
+        shape = shape_node.get_tensor_value()
+        shape = np.array(list(shape), dtype=np.int64)
+        onnx_tensor = numpy_helper.from_array(shape, name)
+        ctx._initializers[name] = onnx_tensor
+        shape_node.set_attr("value", onnx_tensor)
+        return [node]
+    else:
+        op_name = utils.make_name(node.name)
+        cast_op = ctx.insert_new_node_on_input(node, "Cast", name, name=op_name)
+        cast_op.set_attr("to", onnx_pb.TensorProto.INT64)
+        ctx.copy_shape(name, op_name + ":0")
+        return [cast_op, node]
 
+# pylint: disable=W0613,C0111,W0612
 
 def no_op(ctx, node, name, args):
     """Skip node."""
@@ -255,23 +272,8 @@ def reshape_op(ctx, node, name, args):
 
 
 def reshape_op5(ctx, node, name, args):
-    shape_node = node.inputs[1]
     # onnx wants reshape.input[1] to have the value be int64 which is not the case for tensorflow.
-    name = node.input[1]
-    if shape_node.is_const():
-        # if it is a const, change the const to be int64
-        shape = shape_node.get_tensor_value()
-        shape = np.array(list(shape), dtype=np.int64)
-        onnx_tensor = numpy_helper.from_array(shape, name)
-        ctx._initializers[name] = onnx_tensor
-        shape_node.set_attr("value", onnx_tensor)
-        return node
-    else:
-        op_name = utils.make_name(node.name)
-        cast_op = ctx.insert_new_node_on_input(node, "Cast", name, name=op_name)
-        cast_op.set_attr("to", onnx_pb.TensorProto.INT64)
-        ctx.copy_shape(name, op_name + ":0")
-        return [cast_op, node]
+    return _convert_shapenode_to_int64(ctx, node, 1)
 
 
 NCHW_TO_NHWC = [0, 2, 3, 1]
@@ -690,7 +692,7 @@ def expanddims_op(ctx, node, name, args):
 def expanddims_op7(ctx, node, name, args):
     shape = ctx.get_shape(node.output[0])
     shape_name = utils.make_name(node.name)
-    shape_node = ctx.make_const(shape_name, "Const", np.array(shape))
+    shape_node = ctx.make_const(shape_name, "Const", np.array(shape, dtype=np.int64))
     node.type = "Reshape"
     node.input[1] = shape_name
     return node
@@ -786,6 +788,11 @@ def topk_op(ctx, node, name, args):
     return node
 
 
+def tile_op7(ctx, node, name, args):
+    # onnx wants shape input to be int64
+    return _convert_shapenode_to_int64(ctx, node, 1)
+
+
 # pylint: enable=W0613,C0111,W0612
 
 # map tensorflow ops to onnx ops. The format below is
@@ -881,7 +888,7 @@ def topk_op(ctx, node, name, args):
 }
 
 _OPSET_7 = {
-    "Tile": (direct_op, []),
+    "Tile": (tile_op7, []),
     "ResizeNearestNeighbor": (upsample_op, []),
     "BiasAdd": (biasadd_op7, []),
     "BiasAddV1": (biasadd_op7, []),