Fix bug with BatchMatMulV2 when it has the same input twice (#1483)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tf2onnx/graph.py

@@ -1258,7 +1258,7 @@ def remove_input(self, node, to_be_removed, input_index=None):
 
         # don't remove output from parent since others might depend on it
 
-    def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=None, **kwargs):
+    def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=None, input_index=None, **kwargs):
         """Create and insert a new node into the graph.
         Args:
             node: we want to replace the input for this node
@@ -1279,10 +1279,13 @@ def insert_new_node_on_input(self, node, op_type, input_name, name=None, domain=
             input_name = [input_name]
 
         new_node = self.make_node(op_type, input_name, attr=kwargs, outputs=[new_output], name=name, domain=domain)
-        for i, n in enumerate(node.input):
-            if n == input_name[0]:
-                self.replace_input(node, node.input[i], new_output, i)
-                break
+        if input_index is None:
+            for i, n in enumerate(node.input):
+                if n == input_name[0]:
+                    self.replace_input(node, node.input[i], new_output, i)
+                    break
+        else:
+            self.replace_input(node, node.input[input_index], new_output, input_index)
         return new_node
 
     def insert_node_on_output(self, node, output_name=None):

tf2onnx/onnx_opset/math.py

@@ -374,7 +374,7 @@ def version_1(cls, ctx, node, **kwargs):
                 tmp = perm[-1]
                 perm[-1] = perm[-2]
                 perm[-2] = tmp
-                ctx.insert_new_node_on_input(node, "Transpose", node.input[0], perm=perm)
+                ctx.insert_new_node_on_input(node, "Transpose", node.input[0], input_index=0, perm=perm)
 
         if transpose_b != 0:
             shape = ctx.get_shape(node.input[1])
@@ -383,7 +383,7 @@ def version_1(cls, ctx, node, **kwargs):
                 tmp = perm[-1]
                 perm[-1] = perm[-2]
                 perm[-2] = tmp
-                ctx.insert_new_node_on_input(node, "Transpose", node.input[1], perm=perm)
+                ctx.insert_new_node_on_input(node, "Transpose", node.input[1], input_index=1, perm=perm)
 
         unsupported = ["a_is_sparse", "b_is_sparse"]
         for i in unsupported: