Merge pull request #1241 from onnx/tom/opset13updates2

More bug fixes for opset 13 support

Author: TomWildenhain-Microsoft

tf2onnx/onnx_opset/controlflow.py

@@ -286,24 +286,17 @@ class TensorListGetItem:
     def version_7(cls, ctx, node, **kwargs):
         ctx.ta_reads.append(node.input[0])
         node.type = "Gather"
-        ctx.replace_inputs(node, [node.input[0], node.input[1]])
-        ctx.insert_new_node_on_input(node, "Unsqueeze", node.input[1], name=node.child_name(), axes=[0])
-        ctx.insert_new_node_on_output("Squeeze", node.output[0], name=node.child_name(), axes=[0])
-
-    @classmethod
-    def version_13(cls, ctx, node, **kwargs):
-        ctx.ta_reads.append(node.input[0])
-        node.type = "Gather"
-        ctx.replace_inputs(node, [node.input[0], node.input[1]])
-
         g = GraphBuilder(ctx)
 
         usq_node = g.make_unsqueeze({"axes": [0], 'data': node.input[1]}, name=node.child_name(), return_node=True)
-        ctx.insert_node_on_output(usq_node)
-
+        ctx.replace_inputs(node, [node.input[0], usq_node.output[0]])
         sq_node = g.make_squeeze({"axes": [0], 'data': node.output[0]}, name=node.child_name(), return_node=True)
         ctx.insert_node_on_output(sq_node)
 
+    @classmethod
+    def version_13(cls, ctx, node, **kwargs):
+        cls.version_7(ctx, node, **kwargs)
+
 
 @tf_op(["TensorListLength"])
 class TensorListLength:

tf2onnx/onnx_opset/nn.py

@@ -353,7 +353,7 @@ def any_version(cls, opset, ctx, node, **kwargs):
         if len(input_shape) == spatial + 1:
             gb = GraphBuilder(ctx)
             usq_node = gb.make_unsqueeze({"axes": [0], 'data': node.input[0]}, return_node=True)
-            ctx.insert_node_on_output(usq_node, node.input[0])
+            ctx.replace_inputs(node, [usq_node.output[0]] + node.input[1:])
 
         # Set padding.
         add_padding(

tf2onnx/onnx_opset/tensor.py

@@ -1256,8 +1256,8 @@ def any_version_after9(cls, opset, ctx, node, **kwargs):
         depth = GraphBuilder(ctx).make_unsqueeze({'data': node.input[1], 'axes': [0]})
         on_value = node.input[2]
         off_value = node.input[3]
-        on_value = ctx.make_node("Unsqueeze", [on_value], attr={"axes": [0]}).output[0]
-        off_value = ctx.make_node("Unsqueeze", [off_value], attr={"axes": [0]}).output[0]
+        on_value = GraphBuilder(ctx).make_unsqueeze({'data': on_value, 'axes': [0]})
+        off_value = GraphBuilder(ctx).make_unsqueeze({'data': off_value, 'axes': [0]})
         off_on_value = ctx.make_node("Concat", [off_value, on_value], attr={"axis": 0}).output[0]
 
         indices = node.input[0]
@@ -2385,15 +2385,17 @@ def normalize():
         pad_length_2 = body_graph.make_node('Concat', [zeo, pad_length.output[0]], attr={'axis': 0})
         padded_range = body_graph.make_node('Pad', [sliced_range.output[0], pad_length_2.output[0]])
         # opset == 11, no need to change unsqueeze
-        unsqueezed_range = body_graph.make_node('Unsqueeze', [padded_range.output[0]], attr={'axes': [1]})
+        unsqueezed_range = GraphBuilder(body_graph).make_unsqueeze(
+            {'data': padded_range.output[0], 'axes': [1]}, return_node=True)
         half_shape_x = body_graph.make_node('Slice',
                                             [new_shape.output[0], zeo, minus_two])
         shape_range = body_graph.make_node('Shape', [unsqueezed_range.output[0]])
         full_shape = body_graph.make_node('Concat', [half_shape_x.output[0], shape_range.output[0]], attr={'axis': 0})
         expanded_range = body_graph.make_node('Expand', [unsqueezed_range.output[0], full_shape.output[0]])
         gathered_input = body_graph.make_node('GatherElements', [processed_input.output[0], expanded_range.output[0]],
                                               attr={'axis': -1})
-        squeezed_input = body_graph.make_node('Squeeze', [gathered_input.output[0]], attr={'axes': [-1]})
+        squeezed_input = GraphBuilder(body_graph).make_squeeze(
+            {'data': gathered_input.output[0], 'axes': [-1]}, return_node=True)
         left_width = body_graph.make_node('Sub', [new_width.output[0], abs_k.output[0]])
         dims = body_graph.make_node('Concat', [left_width.output[0], new_depth.output[0]], attr={'axis': 0})
         valid_dim = body_graph.make_node('ReduceMin', [dims.output[0]])
@@ -2505,8 +2507,8 @@ def normalize():
                                                raw_output_shape + [-1])
         squeeze_sliced_graph = ctx.create_new_graph_with_same_config()
         squeeze_sliced_graph.parent_graph = ctx
-        squeeze_sliced = squeeze_sliced_graph.make_node('Squeeze', [final_output_right_sliced.output[0]],
-                                                        attr={'axes': [-2]})
+        squeeze_sliced = GraphBuilder(squeeze_sliced_graph).make_squeeze(
+            {'data': final_output_right_sliced.output[0], 'axes': [-2]}, return_node=True)
         squeeze_sliced_graph.add_graph_output(squeeze_sliced.output[0], ctx.get_dtype(node.input[0]), raw_output_shape)
         shapes = node.output_shapes
         dtypes = node.output_dtypes
@@ -2680,14 +2682,14 @@ def version_13(cls, ctx, node, **kwargs):
 @tf_op(["MatrixDiag", "MatrixDiagV2", "MatrixDiagV3"])
 class MatrixDiag:
     @classmethod
-    def any_version(cls, opset, ctx, node, **kwargs):
+    def version_12(cls, ctx, node, **kwargs):
         # Assemble MatrixDiagV3 by ReverseSequence
         argc = len(node.input)
 
-        if opset >= 13:
-            squeeze_axes0 = ctx.make_const(utils.make_name("const_axes"), np.array([0], dtype=np.int64))
-            squeeze_axes_1 = ctx.make_const(utils.make_name("const_axes"), np.array([-1], dtype=np.int64))
-            squeeze_axes_2 = ctx.make_const(utils.make_name("const_axes"), np.array([-2], dtype=np.int64))
+        if ctx.opset >= 13:
+            squeeze_axes0 = ctx.make_const(utils.make_name("const_axes"), np.array([0], dtype=np.int64)).output[0]
+            squeeze_axes_1 = ctx.make_const(utils.make_name("const_axes"), np.array([-1], dtype=np.int64)).output[0]
+            squeeze_axes_2 = ctx.make_const(utils.make_name("const_axes"), np.array([-2], dtype=np.int64)).output[0]
 
         minus_two, minus_one, zeo, one, two = \
             [n.output[0] for n in ctx.make_consts([[-2], [-1], [0], [1], [2]])]
@@ -2712,7 +2714,7 @@ def processdiag():
             diag = node.input[0]
             shape = ctx.get_shape(diag)
             if len(shape) == 1:
-                if opset < 13:
+                if ctx.opset < 13:
                     diag = mknode("Unsqueeze", [diag], attr={"axes": [0]})
                 else:
                     diag = mknode("Unsqueeze", [diag, squeeze_axes0])
@@ -2737,7 +2739,7 @@ def id_diag():
             def ex_diag():
                 g = ctx.create_new_graph_with_same_config()
                 g.parent_graph = ctx
-                if opset < 13:
+                if ctx.opset < 13:
                     ex = mknode2(g, "Unsqueeze", [diag], attr={"axes": [-2]})
                 else:
                     ex = mknode2(g, "Unsqueeze", [diag, squeeze_axes_2])
@@ -2755,7 +2757,7 @@ def squeeze_12(name):
         def squeeze_13(name):
             return ctx.make_node("Squeeze", [name, squeeze_axes_1]).output[0]
 
-        squeeze = squeeze_12 if opset < 13 else squeeze_13
+        squeeze = squeeze_12 if ctx.opset < 13 else squeeze_13
 
         # gather inputs
         diag, k, k_min, k_max, k_max_nxt = processdiag()
@@ -3018,14 +3020,10 @@ def paddiag():
         ctx.make_node("Identity", [padded], name=node.name,
                       outputs=node.output, shapes=shapes, dtypes=dtypes)
 
-    @classmethod
-    def version_12(cls, ctx, node, **kwargs):
-        cls.any_version(12, ctx, node, **kwargs)
-
     @classmethod
     def version_13(cls, ctx, node, **kwargs):
         # Parameters moved to inputs for operator Squeeze, Unsqueeze.
-        cls.any_version(13, ctx, node, **kwargs)
+        cls.version_12(ctx, node, **kwargs)
 
 
 @tf_op("MatrixSetDiagV3")

tf2onnx/rewriter/gru_rewriter.py

@@ -211,7 +211,8 @@ def process_var_init_nodes(self, context):
             const_node = self.g.make_const(initial_name, new_val)
             context.onnx_input_ids["initial_state"] = const_node.output[0]
             return
-        squeeze_node = self.g.make_node("Unsqueeze", [initializer_input_id], attr={"axes": [0]})
+        squeeze_node = GraphBuilder(self.g).make_unsqueeze(
+            {'data': initializer_input_id, 'axes': [0]}, return_node=True)
         to_replace = [n for n in self.g.get_nodes() if n != squeeze_node]
         self.g.replace_all_inputs(initializer_input_id, squeeze_node.output[0], ops=to_replace)
         context.onnx_input_ids["initial_state"] = squeeze_node.output[0]