Tom/enable transpose any rank (#1453)

* Check shapes and dtypes in optimizer tests

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

* Enable transpose optimization of any rank

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

* Fix shape bug in transpose opt

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

Author: TomWildenhain-Microsoft

tests/run_pretrained_models.yaml

@@ -552,7 +552,7 @@ mobilebert_tflite:
 
 palm_detection_tflite:
   tf_min_version: 2.1
-  disabled: false   # Converts but produces incorrect results. Working on fixing it.
+  disabled: false
   url: https://github.com/google/mediapipe/raw/master/mediapipe/modules/palm_detection/palm_detection.tflite
   model: "palm_detection.tflite"
   model_type: tflite
@@ -568,3 +568,21 @@ palm_detection_tflite:
     - classificators
   rtol: 0.02
   atol: 0.0005
+
+melgan_tflite:   # TFLite model with FlexOps and rank-3 transposes
+  tf_min_version: 2.1
+  disabled: false
+  url: https://tfhub.dev/tulasiram58827/lite-model/melgan/dr/1?lite-format=tflite
+  model: "melgan.tflite"
+  model_type: tflite
+  input_get: get_zeros
+  opset_constraints:
+    "onnx":
+      "min": 11
+  dequantize: true
+  inputs:
+    "mels": [1, 100, 80]
+  outputs:
+    - Identity
+  rtol: 0.02
+  atol: 0.0005

tests/test_optimizers.py

@@ -27,8 +27,10 @@ def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
                         remaining_op_num, debug=False, rtol=1e-07):
         utils.make_sure(op_type is not None, "op_type should be specified")
         utils.make_sure(remaining_op_num is not None, "remaining_op_num should be specified")
+        utils.make_sure(self.config.is_onnxruntime_backend, "only onnxruntime is supported to test transpose optimizer")
 
         origin_model_path = self.save_onnx_model(origin_proto, onnx_feed_dict, postfix="_origin")
+        expected = self.run_onnxruntime(origin_model_path, onnx_feed_dict, output_names_with_port)
 
         new_proto, new_graph = GraphUtil.optimize_model_proto(origin_proto, catch_errors=False, return_graph=True)
 
@@ -37,21 +39,16 @@ def run_and_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
         new_model_path = self.save_onnx_model(new_proto, onnx_feed_dict, postfix="_opt")
         current = GraphUtil.get_node_count_from_onnx_graph(new_proto.graph)
 
-        self.assertTrue(current[op_type] == remaining_op_num,
-                        msg="Expect " + str(remaining_op_num) + " " + op_type + " ops left, but actually " + str(
-                            current[op_type]) + " left")
-
-        if self.config.is_onnxruntime_backend:
-            expected = self.run_onnxruntime(origin_model_path, onnx_feed_dict, output_names_with_port)
-            actual = self.run_onnxruntime(new_model_path, onnx_feed_dict, output_names_with_port)
-        else:
-            raise ValueError("only onnxruntime is supported to test transpose optimizer")
+        actual = self.run_onnxruntime(new_model_path, onnx_feed_dict, output_names_with_port)
 
         for expected_val, actual_val in zip(expected, actual):
             self.assertAllClose(expected_val, actual_val, rtol=rtol, atol=1e-5)
             self.assertEqual(expected_val.dtype, actual_val.dtype)
             self.assertEqual(expected_val.shape, actual_val.shape)
 
+        self.assertTrue(current[op_type] == remaining_op_num,
+                        msg="Expect " + str(remaining_op_num) + " " + op_type + " ops left, but actually " + str(
+                            current[op_type]) + " left")
         self.assert_shapes_correct(new_graph, allow_missing=False, run_checker=True)
 
         return new_proto
@@ -124,6 +121,7 @@ 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]),
     ])
@@ -176,6 +174,7 @@ def test_transpose_with_add2(self, input_shape1, input_shape2, perm_input, 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]),
     ])
@@ -196,6 +195,7 @@ def test_transpose_relu(self, shape, perm_input, perm_output):
                                    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]),
     ])
@@ -216,6 +216,7 @@ def test_transpose_leaky_relu(self, shape, perm_input, perm_output):
                                    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]),
     ])
@@ -240,15 +241,16 @@ def test_transpose_quantize(self, shape, perm_input, perm_output):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), [0, 2, 1], [0, 2, 1]),
         ((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]),
     ])
     @check_opset_min_version(13, "QuantizeLinear with axis")
     def test_transpose_quantize_with_axis(self, shape, perm_input, perm_output):
-        scale = numpy_helper.from_array(np.array([0.75, 0.1, 2.3, 0.3, 0.42], dtype=np.float32), name='scale')
-        zero_point = numpy_helper.from_array(np.array([2, 4, 6, 8, 10], dtype=np.uint8), name='zero_point')
+        scale = numpy_helper.from_array(np.array([0.75, 0.1, 2.3, 0.3], dtype=np.float32), name='scale')
+        zero_point = numpy_helper.from_array(np.array([2, 4, 6, 8], dtype=np.uint8), name='zero_point')
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
-        node2 = helper.make_node("QuantizeLinear", ["Y", "scale", "zero_point"], ["Z"], name="quantize", axis=2)
+        node2 = helper.make_node("QuantizeLinear", ["Y", "scale", "zero_point"], ["Z"], name="quantize", axis=1)
         node3 = helper.make_node("Transpose", ["Z"], ["Z1"], perm=perm_output, name="trans_2")
 
         graph = helper.make_graph(
@@ -264,6 +266,7 @@ def test_transpose_quantize_with_axis(self, shape, perm_input, perm_output):
                                    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]),
     ])
@@ -288,15 +291,16 @@ def test_transpose_dequantize(self, shape, perm_input, perm_output):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), [0, 2, 1], [0, 2, 1]),
         ((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]),
     ])
     @check_opset_min_version(13, "DequantizeLinear with axis")
     def test_transpose_dequantize_with_axis(self, shape, perm_input, perm_output):
-        scale = numpy_helper.from_array(np.array([0.75, 0.1, 2.3, 0.3, 0.42], dtype=np.float32), name='scale')
-        zero_point = numpy_helper.from_array(np.array([2, 4, 6, 8, 10], dtype=np.uint8), name='zero_point')
+        scale = numpy_helper.from_array(np.array([0.75, 0.1, 2.3, 0.3], dtype=np.float32), name='scale')
+        zero_point = numpy_helper.from_array(np.array([2, 4, 6, 8], dtype=np.uint8), name='zero_point')
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
-        node2 = helper.make_node("DequantizeLinear", ["Y", "scale", "zero_point"], ["Z"], name="dequantize", axis=2)
+        node2 = helper.make_node("DequantizeLinear", ["Y", "scale", "zero_point"], ["Z"], name="dequantize", axis=1)
         node3 = helper.make_node("Transpose", ["Z"], ["Z1"], perm=perm_output, name="trans_2")
 
         graph = helper.make_graph(
@@ -312,6 +316,7 @@ def test_transpose_dequantize_with_axis(self, shape, perm_input, perm_output):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ([2, 3, 4], [1, 2, 1], [1], [0, 2, 1], [0, 2, 1]),
         ([2, 3, 4, 5], [1, 2, 1, 2], [1], [0, 2, 3, 1], [0, 3, 1, 2]),
         ([2, 3, 4, 5], [1, 2, 1, 2], [1, 2], [0, 2, 3, 1], [0, 3, 1, 2]),
         ([2, 3, 4, 5], [1, 2, 1, 2], [0, 1, 2, 3], [0, 2, 3, 1], [0, 3, 1, 2]),
@@ -351,6 +356,7 @@ def test_transpose_slice(self, input_shape, slice_size, axes, perm_input, perm_o
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ([2, 3, 4], [1, 2, 1], [1], [0, 2, 1], [0, 2, 1]),
         ([2, 3, 4, 5], [1, 2, 1, 2], [1], [0, 2, 3, 1], [0, 3, 1, 2]),
         ([2, 3, 4, 5], [1, 2, 1, 2], [1, 2], [0, 2, 3, 1], [0, 3, 1, 2]),
         ([2, 3, 4, 5], [1, 2, 1, 2], [0, 1, 2, 3], [0, 2, 3, 1], [0, 3, 1, 2]),
@@ -390,6 +396,7 @@ def test_transpose_slice_opset_10(self, input_shape, slice_size, axes, perm_inpu
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), (4, 2, 3), (2, 0, 1), (1, 2, 0)),
         ((2, 3, 4, 5), (2, 4, 5, 3), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((2, 3, 4, 5, 6), (2, 4, 5, 6, 3), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -513,6 +520,7 @@ def test_transpose_merge(self, input_shape1, input_shape2, perm):
 
 
     @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]),
     ])
@@ -533,6 +541,7 @@ def test_transpose_mul_as_square(self, shape, perm_input, perm_output):
                                    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]),
     ])
@@ -555,6 +564,7 @@ def test_transpose_mul_broadcastable_const(self, shape, perm_input, perm_output)
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), [2, 0, 1]),
         ((2, 3, 4, 5), [0, 2, 3, 1]),
         ((2, 3, 4, 5, 6), [0, 2, 3, 4, 1]),
     ])
@@ -574,6 +584,7 @@ def test_transpose_with_shape(self, shape, perm):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), (4, 2, 3), [2, 0, 1]),
         ((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]),
     ])
@@ -593,6 +604,7 @@ def test_transpose_with_identity(self, input_shape, output_shape, perm):
                                    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]),
     ])
@@ -613,6 +625,7 @@ def test_transpose_sqrt(self, shape, perm_input, perm_output):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((1, 3, 4), [4, 3], [0, 2, 1], [1, 0]),
         ((1, 3, 4, 5), (4, 5, 3), [0, 2, 3, 1], [1, 2, 0]),
         ((1, 3, 4, 5, 6), (4, 5, 6, 3), [0, 2, 3, 4, 1], [1, 2, 3, 0]),
     ])
@@ -635,17 +648,18 @@ def test_transpose_with_squeeze1(self, input_shape, output_shape, perm, expected
         self.check_transpose_perm(model_after_opt, expected_perm)
 
     @parameterized.expand([
-        ((1, 3, 4, 5), (1, 1, 4, 5, 1, 3, 1), [0, 2, 3, 1], [0, 1, 4, 5, 2, 3, 6]),
-        ((1, 3, 4, 5, 6), (1, 1, 4, 5, 1, 6, 1, 3), [0, 2, 3, 4, 1], [0, 1, 4, 5, 6, 7, 2, 3]),
+        ((1, 3, 4), (1, 4, 1, 3, 1, 1), [2, 0, 1], [0, 4, 5], [2, 3, 0, 1, 4, 5]),
+        ((1, 3, 4, 5), (1, 1, 4, 5, 1, 3, 1), [0, 2, 3, 1], [0, 4, 6], [0, 1, 4, 5, 2, 3, 6]),
+        ((1, 3, 4, 5, 6), (1, 1, 4, 5, 1, 6, 1, 3), [0, 2, 3, 4, 1], [0, 4, 6], [0, 1, 4, 5, 6, 7, 2, 3]),
     ])
-    def test_transpose_with_unsqueeze(self, input_shape, output_shape, perm, expected_perm):
+    def test_transpose_with_unsqueeze(self, input_shape, output_shape, perm, axes_val, expected_perm):
         # unsqueeze the first dim
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm, name="trans")
         if self.config.opset <= 12:
-            node2 = helper.make_node("Unsqueeze", ["Y"], ["Z"], name="unsqueeze", axes=[0, 4, 6])
+            node2 = helper.make_node("Unsqueeze", ["Y"], ["Z"], name="unsqueeze", axes=axes_val)
             nodes = [node1, node2]
         else:
-            axes = self._make_onnx_const(np.array([0, 4, 6], dtype=np.int64), "axes")
+            axes = self._make_onnx_const(np.array(axes_val, dtype=np.int64), "axes")
             node2 = helper.make_node("Unsqueeze", ["Y", "axes"], ["Z"], name="unsqueeze")
             nodes = [axes, node1, node2]
 
@@ -662,6 +676,7 @@ def test_transpose_with_unsqueeze(self, input_shape, output_shape, perm, expecte
         self.check_transpose_perm(model_after_opt, expected_perm)
 
     @parameterized.expand([
+        ((1, 3, 4), [4, 3], [0, 2, 1], [1, 0]),
         ((1, 3, 4, 5), (4, 5, 3), [0, 2, 3, 1], [1, 2, 0]),
         ((1, 3, 4, 5, 6), (4, 5, 6, 3), [0, 2, 3, 4, 1], [1, 2, 3, 0]),
     ])
@@ -816,6 +831,7 @@ def test_transpose_with_squeeze4_13(self, input_shape, output_shape, perm):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((10, 3, 4), [0, 2, 1], [0, 2, 1]),
         ((10, 3, 4, 5), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((10, 3, 4, 5, 6), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -880,6 +896,7 @@ def _make_loop(external_inputs, outputs):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), [4, 2, 3], [2, 0, 1], [1, 2, 0]),
         ((2, 3, 4, 5), [2, 4, 5, 3], [0, 2, 3, 1], [0, 3, 1, 2]),
         ((2, 3, 4, 5, 6), [2, 4, 5, 6, 3], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -965,6 +982,7 @@ def test_transpose_add_with_input_non_const(self, input_shape1, input_shape2, pe
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), [4, 2, 3], [2, 0, 1], [1, 2, 0]),
         ((1, 1, 3, 3), (1, 3, 3, 1), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 1, 3, 3, 3), (1, 3, 3, 3, 1), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -1050,6 +1068,7 @@ def test_transpose_add_with_conv_2(self, input_shape, weights_shape, output_shap
                                    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]),
         ((1, 3, 4, 5, 6), (2, 5, 6, 8, 10), [1, 0, 1, 3, 1, 0, 2, 2, 1, 1], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -1071,6 +1090,7 @@ def test_transpose_pad(self, input_shape, output_shape, pads, perm_input, perm_o
                                    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]),
         ((1, 3, 4, 5, 6), (2, 5, 6, 8, 10), [1, 0, 1, 3, 1, 0, 2, 2, 1, 1], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -1097,6 +1117,7 @@ def test_transpose_pad11(self, input_shape, output_shape, pads, perm_input, perm
                                    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]),
         ((1, 3, 4, 5, 6), (2, 5, 6, 8, 10), [1, 0, 1, 3, 1, 0, 2, 2, 1, 1], [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -1125,6 +1146,7 @@ def test_transpose_pad11_non_const_pads(self, input_shape, output_shape, pads, p
                                    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]),
     ])
@@ -1145,6 +1167,7 @@ def test_transpose_reciprocal(self, shape, perm_input, perm_output):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((3, 4, 5), (3, 4, 1), [0, 2, 1], [0, 2, 1]),
         ((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]),
     ])
@@ -1166,6 +1189,7 @@ def test_transpose_reducemean(self, input_shape, output_shape, perm_input, perm_
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((3, 4, 5), (3, 4, 1), [1], [0, 2, 1], [0, 2, 1]),
         ((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]),
@@ -1258,6 +1282,7 @@ def test_transpose_tile(self):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((3, 4, 5), (3, 4, 1), [1], [0, 2, 1], [0, 2, 1]),
         ((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]),
@@ -1286,6 +1311,7 @@ def test_transpose_reducesum_opset_13(self, input_shape, output_shape, axes, per
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((2, 3, 4), (4, 2, 3), [2, 0, 1]),
         ((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]),
     ])
@@ -1362,6 +1388,7 @@ def test_trans_can_be_replaced_with_reshape2(self, input_shape_np, input_shape,
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((1, 6, 8), [2, 0, 1], [1, 2, 0]),
         ((1, 6, 8, 9), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 6, 8, 9, 2), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])
@@ -1388,6 +1415,7 @@ def test_two_transposes_switch_with_mul(self, shape, perm_input, perm_output):
                                    model_proto, remaining_transpose_num=0)
 
     @parameterized.expand([
+        ((1, 6, 8), (8, 1, 6), [2, 0, 1], [1, 2, 0]),
         ((1, 6, 8, 9), (1, 8, 9, 6), [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 6, 8, 9, 2), (1, 8, 9, 2, 6), [0, 2, 3, 4, 1], [0, 4, 1, 2, 3]),
     ])