Handle Mul as Square in TransposeOptimizer (#1339)

Signed-off-by: Mateusz Tabaka <[email protected]>

Co-authored-by: TomWildenhain-Microsoft <[email protected]>

Author: mateusztabaka

tests/test_optimizers.py

@@ -477,6 +477,28 @@ def test_transpose_merge(self, input_shape1, input_shape2, perm):
         self.run_transpose_compare(["OUT"], {"X": np.random.randn(*input_shape1).astype(np.float32)},
                                    model_proto, remaining_transpose_num=1)
 
+
+    @parameterized.expand([
+        ((2, 3, 4, 5), [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((2, 3, 4, 5, 6), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+    ])
+    def test_transpose_mul_as_square(self, shape, perm_input, perm_output):
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans")
+        node1 = helper.make_node("Mul", ["Y", "Y"], ["Z"], name="mul")
+        node2 = helper.make_node("Transpose", ["Z"], ["OUT"], perm=perm_output, name="trans_1")
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-mul-as-sqr-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, shape)],
+            [helper.make_tensor_value_info("OUT", TensorProto.FLOAT, shape)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["OUT"], {"X": np.random.randn(*shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
+
     @parameterized.expand([
         ((2, 3, 4, 5), [0, 2, 3, 1]),
         ((2, 3, 4, 5, 6), [0, 2, 3, 4, 1]),

tf2onnx/optimizer/transpose_optimizer.py

@@ -460,7 +460,13 @@ def _mul_handler(self, trans, node):
 
         # node's inputs may come from one same node. if so the multiplier_input_node may be none
         if multiplier_input_node is None:
-            return False
+            if not self._nodes_has_single_consumer_node([trans]):
+                return False
+            self._g.replace_all_inputs(node.output[0], trans.output[0])
+            self._g.replace_input(node, node.input[0], trans.input[0], 0)
+            self._g.replace_input(node, node.input[1], trans.input[0], 1)
+            self._g.replace_input(trans, trans.input[0], node.output[0], 0)
+            return True
 
         # convert  mul(trans(x), trans(y)) ->  trans(mul(x, y))
         if multiplier_input_node.type == "Transpose":