Handle Abs in TransposeOptimizer (#1699)

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

Author: mateusztabaka

tests/test_optimizers.py

@@ -116,6 +116,27 @@ def test_transpose_with_concat(self, input_shape, perm, inner_perm):
                          }
             self.run_transpose_compare(["res"], feed_dict, model_proto, remaining_transpose_num=1)
 
+    @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_abs(self, shape, perm_input, perm_output):
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans1")
+        node1 = helper.make_node("Abs", ["Y"], ["Z"], name="abs")
+        node2 = helper.make_node("Transpose", ["Z"], ["OUT"], perm=perm_output, name="trans2")
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-abs-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), [2, 0, 1], [1, 2, 0]),
         ((2, 3, 4, 5), [0, 2, 3, 1], [0, 3, 1, 2]),

tf2onnx/optimizer/transpose_optimizer.py

@@ -193,6 +193,7 @@ def _optimize_at_current_graph_level(self, graph):
 
     def _initialize_handlers(self):
         self._handler_map = {
+            "Abs": self._simple_through_handler,
             "Add": self._add_handler,
             "ArgMax": self._arg_min_max_handler,
             "ArgMin": self._arg_min_max_handler,