refactor

Author: zhijxu-MS

tests/test_optimizers.py

@@ -147,6 +147,21 @@ def test_transpose_with_shape(self):
         self.run_transpose_compare(["Z"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    def test_transpose_with_identity(self):
+        node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans")
+        node2 = helper.make_node("Identity", ["Y"], ["Z"], name="identity")
+
+        graph = helper.make_graph(
+            [node1, node2],
+            "transpose_with_identity",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4, 5))],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 4, 5, 3))],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["Z"], {"X": np.random.randn(2, 3, 4, 5).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=1)
+
     # Tranpose Optimizer Tests End
 
     # Identity Optimizer Tests Start

tf2onnx/optimizer/transpose_optimizer.py

@@ -385,6 +385,7 @@ def _mul_handler(self, trans, node):
                 multiplier_input_id = i
                 multiplier_input_node = input_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 or not multiplier_input_node.is_const():
             return False
         multiplier = multiplier_input_node.get_tensor_value(as_list=False)