Add PRelu to Transpose optimizer (#1630)

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

Co-authored-by: Guenther Schmuelling <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_optimizers.py

@@ -211,6 +211,31 @@ def test_transpose_leaky_relu(self, shape, perm_input, perm_output):
         self.run_transpose_compare(["Z1"], {"X": np.random.randn(*shape).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @parameterized.expand([
+        ((2, 3, 4), [2, 0, 1], [1, 2, 0]),
+        ((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_with_prelu(self, input_shape, perm_input, perm_output):
+        node1 = helper.make_node("Transpose", ["input_data1"], ["Y"], perm=perm_input, name="trans")
+        node2 = helper.make_node("PRelu", ["Y", "input_data2"], ["Z"], name="add")
+        node3 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm_output, name="trans2")
+
+        graph = helper.make_graph(
+            [node1, node2, node3],
+            "transpose_with_shape",
+            [helper.make_tensor_value_info("input_data1", TensorProto.FLOAT, input_shape),
+             helper.make_tensor_value_info("input_data2", TensorProto.FLOAT, (input_shape[1],)),
+             ],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, input_shape)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        feed_dict = {"input_data1": np.random.randn(*input_shape).astype(np.float32),
+                     "input_data2": np.random.randn(input_shape[1]).astype(np.float32),
+                     }
+        self.run_transpose_compare(["res"], feed_dict, model_proto, remaining_transpose_num=0)
+
     @parameterized.expand([
         ((2, 3, 4), [2, 0, 1], [1, 2, 0]),
         ((2, 3, 4, 5), [0, 2, 3, 1], [0, 3, 1, 2]),

tf2onnx/optimizer/transpose_optimizer.py

@@ -208,6 +208,7 @@ def _initialize_handlers(self):
             "Min": self._maxmin_handler,
             "Mul": self._mul_handler,
             "Pad": self._pad_handler,
+            "PRelu": self._prelu_handler,
             "Reciprocal": self._simple_through_handler,
             "ReduceLogSum": self._reduce_handler,
             "ReduceLogSumExp": self._reduce_handler,
@@ -816,6 +817,9 @@ def permute_pads(pads):
         self._g.replace_input(node, node.input[1], new_pads.output[0], 1)
         return self._switch_transpose_and_node(node, trans)
 
+    def _prelu_handler(self, trans, node):
+        return self._handle_node_having_branches(trans, node)
+
     def _arg_min_max_handler(self, trans, node):
         axis = node.get_attr_value("axis", 0)
         node.set_attr("axes", [axis])