Handle 'Neg' nodes in transpose optimzier

Signed-off-by: Jan Haug <jan@recogni.com>

Author: jan-haug

tests/test_optimizers.py

@@ -1109,6 +1109,27 @@ def test_transpose_add_with_conv_2(self, input_shape, weights_shape, output_shap
  "W": np.random.randn(*weights_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_neg(self, shape, perm_input, perm_output):
+ node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans1")
+ node1 = helper.make_node("Neg", ["Y"], ["Z"], name="neg")
+ node2 = helper.make_node("Transpose", ["Z"], ["OUT"], perm=perm_output, name="trans2")
+
+ graph = helper.make_graph(
+ [node0, node1, node2],
+ "transpose-neg-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([
  ((3, 4, 5), (8, 4, 6), [1, 3, 0, 0, 2, 0], [2, 0, 1], [1, 2, 0]),
  ((1, 3, 4, 5), (2, 6, 4, 8), [1, 0, 1, 3, 0, 0, 2, 0], [0, 2, 3, 1], [0, 3, 1, 2]),

tf2onnx/optimizer/transpose_optimizer.py

@@ -208,6 +208,7 @@ def _initialize_handlers(self):
  "Max": self._maxmin_handler,
  "Min": self._maxmin_handler,
  "Mul": self._mul_handler,
+ "Neg": self._simple_through_handler,
  "Pad": self._pad_handler,
  "PRelu": self._prelu_handler,
  "Reciprocal": self._simple_through_handler,