Handle ReduceSum in TransposeOptimizer (#1342)

* Handle ReduceSum in TransposeOptimizer

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

* Remove unused variable

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

Author: mateusztabaka

tests/test_optimizers.py

@@ -1062,7 +1062,6 @@ def test_transpose_reciprocal(self, shape, perm_input, perm_output):
         self.run_transpose_compare(["OUT"], {"X": np.random.randn(*shape).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
-
     @parameterized.expand([
         ((1, 3, 4, 5), (1, 3, 1, 1), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 3, 4, 5, 6), (1, 3, 1, 1, 1), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
@@ -1084,6 +1083,60 @@ def test_transpose_reducemean(self, input_shape, output_shape, perm_input, perm_
         self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @parameterized.expand([
+        ((1, 3, 4, 5), (1, 3, 4, 1), [2], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5), (1, 3, 1, 1), [1, 2], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5), (1, 1, 1, 1), [0, 1, 2, 3], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5, 6), (1, 3, 1, 5, 6), [1], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+        ((1, 3, 4, 5, 6), (1, 3, 1, 1, 1), [1, 2, 3], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+        ((1, 3, 4, 5, 6), (1, 1, 1, 1, 1), [0, 1, 2, 3, 4], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+    ])
+    @check_opset_max_version(12, "ReduceSum from opset <= 12 has axes as an attribute")
+    def test_transpose_reducesum(self, input_shape, output_shape, axes, perm_input, perm_output):
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
+        node1 = helper.make_node("ReduceSum", ["Y"], ["Z"], axes=axes,
+                                 keepdims=1, name="reducesum")
+        node2 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm_output, name="trans_2")
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-reducesum-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
+    @parameterized.expand([
+        ((1, 3, 4, 5), (1, 3, 4, 1), [2], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5), (1, 3, 1, 1), [1, 2], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5), (1, 1, 1, 1), [0, 1, 2, 3], [0, 2, 3, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5, 6), (1, 3, 1, 5, 6), [1], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+        ((1, 3, 4, 5, 6), (1, 3, 1, 1, 1), [1, 2, 3], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+        ((1, 3, 4, 5, 6), (1, 1, 1, 1, 1), [0, 1, 2, 3, 4], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
+    ])
+    @check_opset_min_version(13, "ReduceSum from opset >= 13 has axes as an input")
+    def test_transpose_reducesum_opset_13(self, input_shape, output_shape, axes, perm_input, perm_output):
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
+        node1 = helper.make_node("ReduceSum", ["Y", "axes"], ["Z"], keepdims=1, name="reducesum")
+        node2 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm_output, name="trans_2")
+
+        axes = np.array(axes, dtype=np.int64)
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-reducesum-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
+            [helper.make_tensor("axes", TensorProto.INT64, axes.shape, axes)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
     @parameterized.expand([
         ((2, 3, 4, 5), (2, 4, 5, 3), [0, 2, 3, 1]),
         ((2, 3, 4, 5, 6), (2, 4, 5, 6, 3), [0, 2, 3, 4, 1]),

tf2onnx/optimizer/transpose_optimizer.py

@@ -193,6 +193,7 @@ def _initialize_handlers(self):
             "Pad": self._pad_handler,
             "Reciprocal": self._simple_through_handler,
             "ReduceMean": self._reducemean_handler,
+            "ReduceSum": self._reducesum_handler,
             "Relu": self._simple_through_handler,
             "Shape": self._shape_handler,
             "Sigmoid": self._simple_through_handler,
@@ -712,6 +713,34 @@ def _reducemean_handler(self, trans, node):
             return self._switch_transpose_and_node(node, trans)
         return False
 
+    def _reducesum_handler(self, trans, node):
+        keepdims = node.get_attr("keepdims")
+        # make sure keepdims is 1, then we can do the swap, otherwise, please don't, because
+        # once keepdims is not set, original dims are lost, so transpose back won't work well.
+        # by default, if keepdims is not specified, it is 1
+        if keepdims and keepdims.i == 0:
+            return False
+        if self._g.opset <= 12:
+            axes = node.get_attr("axes").ints
+            perm = trans.get_attr('perm').ints
+            new_axes = [perm[axis] for axis in axes]
+            node.set_attr("axes", new_axes)
+            return self._switch_transpose_and_node(node, trans)
+        if node.inputs[1].is_const():
+            axes = node.inputs[1].get_tensor_value()
+            perm = trans.get_attr('perm').ints
+            axes = [perm[axes[i]] for i in range(len(axes))]
+            new_axes = np.array(axes, dtype=np.int64)
+            if self._nodes_has_single_consumer_node([node.inputs[1]]):
+                node.inputs[1].set_tensor_value(new_axes)
+            else:
+                new_axes_const = self._g.make_const(
+                    utils.make_name(node.inputs[1].name), new_axes
+                )
+                self._g.replace_input(node, node.input[1], new_axes_const.output[0], 1)
+            return self._switch_transpose_and_node(node, trans)
+        return False
+
     def _slice_handler(self, trans, node):
         trans_rank = get_transpose_rank(trans)
         axes = None