#  Copyright (c) 2020, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause

import unittest

import numpy as np
import pytest

from coremltools.converters.mil.mil import Builder as mb
from coremltools.converters.mil.mil import get_new_symbol
from coremltools.converters.mil.mil.passes.defs.optimize_repeat_ops import TransformAxisUpdateOps
from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
from coremltools.converters.mil.testing_utils import (
    apply_pass_and_basic_check,
    assert_model_is_valid,
    get_op_types_in_program,
)

np.random.seed(1984)


class TransposeOptimizationPass(unittest.TestCase):
    """
    Input graph:
    input -----> transpose(axis=[1,0]) -----> transpose(axis=[1,0]) ---> out

    Output graph:
    input -----> identity -----> out
    """

    def test_simple_consecutive_ops_fusion_direct_output(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
        def prog(x):
            x = mb.transpose(x=x, perm=[1, 0])
            x = mb.transpose(x=x, perm=[1, 0])
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(get_op_types_in_program(prev_prog), ["transpose", "transpose"])
        self.assertEqual(get_op_types_in_program(prog), ["identity"])
        assert_model_is_valid(
            prog,
            {"x": (10, 20)},
            expected_output_shapes={block.outputs[0].name: (10, 20)},
        )

    """
    Input graph:
    input -----> transpose(axis=[1,0]) -----> transpose(axis=[1,0]) ----> relu ---> out

    Output graph:
    input -----> relu -----> out
    """

    def test_simple_consecutive_ops_fusion(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20))])
        def prog(x):
            x = mb.transpose(x=x, perm=[1, 0])
            x = mb.transpose(x=x, perm=[1, 0])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(get_op_types_in_program(prev_prog), ["transpose", "transpose", "relu"])
        self.assertEqual(get_op_types_in_program(prog), ["relu"])
        assert_model_is_valid(
            prog,
            {"x": (10, 20)},
            expected_output_shapes={block.outputs[0].name: (10, 20)},
        )

    """
    Input graph:
    input---->transpose(axis=[0,3,1,2])---->relu---->log--->transpose(axis=[0,2,3,1])--->relu--->out

    Output graph:
    input----->relu----->log----->relu--->out
    """

    def test_linear_graph_two_op_fusion(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 3, 4))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.relu(x=x)
            x = mb.log(x=x)
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "log", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "log", "relu"])
        assert_model_is_valid(
            prog,
            {"x": (1, 2, 3, 4)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 3, 4)},
        )

    """
    Input graph:
    input---->transpose(axis=[0,3,1,2])---->relu---->identity--->transpose(axis=[0,2,3,1])--->relu--->out

    Output graph:
    input----->relu----->identity----->relu--->out
    """

    def test_linear_graph_two_op_fusion_1(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 3, 4))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.relu(x=x)
            x = mb.identity(x=x)
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "identity", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "identity", "relu"])
        assert_model_is_valid(
            prog,
            {"x": (1, 2, 3, 4)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 3, 4)},
        )

    """
    Input graph:
    input(shape=1,2,3,4)---->transpose(axis=[0,3,1,2])---->relu---->log--->transpose(axis=[0,2,3,1])--->relu--->out1(shape=1,2,3,4)
                                                                |
                                                                v
                                                        out2(shape=1,4,2,3)

    Output graph:
    input(shape=1,2,3,4)---->relu---->log--->relu--->out1(shape=1,2,3,4)
                                  |
                                  |----->transpose(axis=[0,3,1,2])----->out2(shape=1,4,2,3)
    """

    def test_fusion_with_output_edge_inbetween(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 3, 4))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x1 = mb.relu(x=x)
            x2 = mb.log(x=x1)
            x3 = mb.transpose(x=x2, perm=[0, 2, 3, 1])
            x4 = mb.relu(x=x3)
            return x4, x1

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "log", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "log", "relu", "transpose"])
        assert_model_is_valid(
            prog,
            {"x": (1, 2, 3, 4)},
            expected_output_shapes={
                block.outputs[0].name: (1, 2, 3, 4),
                block.outputs[1].name: (1, 4, 2, 3),
            },
        )

    """
    Input graph:
    input---->transpose(axis=[0,3,1,2])---->relu---->transpose(axis=[0,2,3,1])--->out

    Output graph:
    input----->relu----->out
    """

    def test_linear_graph_two_op_fusion_with_last_op_removal(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 3, 4))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.relu(x=x)
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(get_op_types_in_program(prev_prog), ["transpose", "relu", "transpose"])
        self.assertEqual(get_op_types_in_program(prog), ["relu"])
        assert_model_is_valid(
            prog,
            {"x": (1, 2, 3, 4)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 3, 4)},
        )

    """
    Input graph:
    input(shape=10,2,3)--->transpose(axis=[0,2,1])----->relu---->transpose(axis=[0,2,1])---->out1
                                                    |
                                                    |
                                                    --->relu----->log---->transpose(axis=[0,2,1])---->out2

    Output graph:
    input(shape=10,2,3)----->relu---->out1
                        |
                        |
                        --->relu----->log---->out2
    """

    def test_multiple_fusions(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 2, 1])
            x1 = mb.relu(x=x)
            x2 = mb.relu(x=x)
            y1 = mb.transpose(x=x1, perm=[0, 2, 1])
            x3 = mb.log(x=x2)
            y2 = mb.transpose(x=x3, perm=[0, 2, 1])
            return y1, y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "relu", "transpose", "log", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "relu", "log"])

        assert prev_block.inputs["x"] == prev_block.find_ops(op_type="transpose")[0].inputs["x"]
        assert block.find_ops(op_type="log")[0].outputs[0] in block.outputs
        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3)},
            expected_output_shapes={
                block.outputs[0].name: (10, 2, 3),
                block.outputs[1].name: (10, 2, 3),
            },
        )

    """
    Input graph:
    input(shape=10,2,3,5)--->transpose(axis=[0,2,3,1])----->relu---->pool----->out1
                                                       |
                                                       |
                                                       --->relu----->log---->transpose(axis=[0,3,1,2])---->out2


    Output graph:
    input(shape=10,2,3,5)----->relu---->transpose(axis=[0,2,3,1])---->pool----->out1
                           |
                           |
                           --->relu----->log---->out2
    """

    def test_partial_fusion_0(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3, 5))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x1 = mb.relu(x=x)
            x2 = mb.relu(x=x)
            y1 = mb.avg_pool(x=x1, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            x3 = mb.log(x=x2)
            y2 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            return y1, y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "relu", "avg_pool", "log", "transpose"],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["relu", "relu", "transpose", "avg_pool", "log"],
        )

        assert prev_block.inputs["x"] == prev_block.find_ops(op_type="transpose")[0].inputs["x"]
        assert block.find_ops(op_type="log")[0].outputs[0] == block.outputs[1]
        assert (
            block.find_ops(op_type="transpose")[0].outputs[0]
            == block.find_ops(op_type="avg_pool")[0].inputs["x"]
        )
        assert list(block.find_ops(op_type="transpose")[0].perm.val) == [0, 2, 3, 1]
        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3, 5)},
            expected_output_shapes={
                block.outputs[0].name: (10, 3, 5, 2),
                block.outputs[1].name: (10, 2, 3, 5),
            },
        )

    """
    Input graph:
    input(shape=10,2,3,5)--->transpose(axis=[0,2,1,3])----->relu---->transpose(axis=[0,2,1,3])---->out1
                                                        |
                                                        |
                                                        --->pool--->log---->transpose(axis=[0,2,1,3])---->out2

    Output graph:
    input(shape=10,2,3,5)----->relu---->out1
                           |
                           |
                           --->transpose(axis=[0,2,1,3])---->pool----->log---->transpose(axis=[0,2,1,3])---->out2
    """

    def test_partial_fusion_1(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3, 5))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 2, 1, 3])
            x1 = mb.relu(x=x)
            x2 = mb.avg_pool(x=x, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            y1 = mb.transpose(x=x1, perm=[0, 2, 1, 3])
            x3 = mb.log(x=x2)
            y2 = mb.transpose(x=x3, perm=[0, 2, 1, 3])
            return y1, y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "avg_pool", "transpose", "log", "transpose"],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["relu", "transpose", "avg_pool", "log", "transpose"],
        )

        assert block.inputs["x"] == block.find_ops(op_type="relu")[0].inputs["x"]
        assert block.outputs[0] == block.find_ops(op_type="relu")[0].outputs[0]
        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3, 5)},
            expected_output_shapes={
                block.outputs[0].name: (10, 2, 3, 5),
                block.outputs[1].name: (10, 2, 3, 5),
            },
        )

    """
    Input graph:

                                                             |-------> transpose(axis=[0,2,1,3]) ---->out1(shape=10,2,3,5)
                                                             |
    input(shape=10,2,3,5)-->relu-->transpose(axis=[0,2,1,3])--->relu--->transpose(axis=[0,2,1,3]) ---->out2(shape=10,2,3,5)
                                                                     |
                                                                     |----->pool--------------->out3(shape=10,3,2,5)
                                                                     |
                                                                     |----->pool--------------->out4(shape=10.3.2.5)


    Output graph:

                        |---->out1(shape=10,2,3,5)
                        |
    input---->relu---------->relu------->out2(shape=10,2,3,5)
                                    |
                                    |----->transpose(axis=[0,2,1,3])--->pool---->out3(shape=10,3,2,5)
                                    |
                                    |----->transpose(axis=[0,2,1,3])---->pool--->out4(shape=10.3.2.5)
    """

    def test_partial_fusion_2(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3, 5))])
        def prog(x):
            x = mb.relu(x=x)
            x = mb.transpose(x=x, perm=[0, 2, 1, 3])
            y1 = mb.transpose(x=x, perm=[0, 2, 1, 3])
            x1 = mb.relu(x=x)
            y2 = mb.transpose(x=x1, perm=[0, 2, 1, 3])
            y3 = mb.avg_pool(x=x1, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            y4 = mb.avg_pool(x=x1, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            return y1, y2, y3, y4

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "relu",
                "transpose",
                "transpose",
                "relu",
                "transpose",
                "avg_pool",
                "avg_pool",
            ],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["relu", "relu", "transpose", "avg_pool", "transpose", "avg_pool"],
        )

        assert block.outputs[0] == block.find_ops(op_type="relu")[0].outputs[0]
        assert block.outputs[1] == block.find_ops(op_type="relu")[1].outputs[0]
        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3, 5)},
            expected_output_shapes={
                # Two consecutive relus are merged, so the first two outputs have the same name. See
                # `test_name_change_depend_on_output` in TestMergeConsecutiveRelus.
                block.outputs[1].name: (10, 2, 3, 5),
                block.outputs[2].name: (10, 3, 2, 5),
                block.outputs[3].name: (10, 3, 2, 5),
            },
            # rdar://100243127 ([PyTorch] Duplicate Output Tensor Doesn't work for neuralnetwork).
            backend=("mlprogram", "fp16"),
        )

    """
    Input graph:

    input(shape=10,2,3,5)-->relu--->transpose(axis=[0,2,1,3])----->transpose(axis=[0,2,1,3])---->out1(shape=10,2,3,5)
                                                               |
                                                               ---->relu------>out2(shape=10,3,2,5)

    Output graph:

    input(shape=10,2,3,5)-->relu---->out1(shape=10,2,3,5)
                                  |
                                   ---->relu--->transpose(axis=[0,2,1,3])------>out2(shape=10,3,2,5)
    """

    def test_partial_fusion_3(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3, 5))])
        def prog(x):
            x = mb.relu(x=x)
            x = mb.transpose(x=x, perm=[0, 2, 1, 3])
            x1 = mb.transpose(x=x, perm=[0, 2, 1, 3])
            x2 = mb.relu(x=x)
            return x1, x2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["relu", "transpose", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "relu", "transpose"])

        assert block.outputs[0] == block.find_ops(op_type="relu")[0].outputs[0]
        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3, 5)},
            expected_output_shapes={
                block.outputs[0].name: (10, 2, 3, 5),
                block.outputs[1].name: (10, 3, 2, 5),
            },
        )

    """
    Input graph:

    input(shape=10,2,3,5)-->relu--->transpose(axis=[0,2,1,3])----->transpose(axis=[0,2,1,3])---->out1(shape=10,2,3,5)
                                                               |
                                                               ------>out2(shape=10,3,2,5)

    Output graph:
    same as input graph as one of the optimizing transpose is connected to model output
    """

    def test_partial_fusion_4(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3, 5))])
        def prog(x):
            x = mb.relu(x=x)
            out2 = mb.transpose(x=x, perm=[0, 2, 1, 3])
            out1 = mb.transpose(x=out2, perm=[0, 2, 1, 3])
            return out1, out2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(get_op_types_in_program(prev_prog), ["relu", "transpose", "transpose"])
        self.assertEqual(get_op_types_in_program(prog), ["relu", "transpose", "transpose"])

        assert block.outputs[1] == block.find_ops(op_type="transpose")[0].outputs[0]
        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3, 5)},
            expected_output_shapes={
                block.outputs[0].name: (10, 2, 3, 5),
                block.outputs[1].name: (10, 3, 2, 5),
            },
        )

    """
    Input graph:
    input(shape=10,2,3,5)-->relu-->transpose(axis=[0,2,1,3])--->relu--->transpose(axis=[0,2,1,3]) ---->out1(shape=10,2,3,5)
                                                                     |
                                                                     |--->relu-->pool--------------->out2(shape=10,3,2,5)
                                                                     |
                                                                     |----->pool--------------->out3(shape=10.3.2.5)


    Output graph:
    same as the input graph as materialization ops are greater than cancel ops
    """

    def test_no_fusion_more_materialization_ops(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3, 5))])
        def prog(x):
            x = mb.relu(x=x)
            x = mb.transpose(x=x, perm=[0, 2, 1, 3])
            x1 = mb.relu(x=x)
            y2 = mb.transpose(x=x1, perm=[0, 2, 1, 3])
            x2 = mb.relu(x=x1)
            y3 = mb.avg_pool(x=x2, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            y4 = mb.avg_pool(x=x1, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            return y2, y3, y4

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["relu", "transpose", "relu", "transpose", "relu", "avg_pool", "avg_pool"],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["relu", "transpose", "relu", "transpose", "relu", "avg_pool", "avg_pool"],
        )

        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3, 5)},
            expected_output_shapes={
                block.outputs[0].name: (10, 2, 3, 5),
                block.outputs[1].name: (10, 3, 2, 5),
                block.outputs[2].name: (10, 3, 2, 5),
            },
        )

    """
    Input graph:
    input(shape=10,2,3)--->transpose(axis=[0,2,1])----->relu---->transpose(axis=[0,2,1])---->out1
                                                    |
                                                    |
                                                    --->reduce(axis=2)----->log---->transpose(axis=[0,2,1])---->out2

    Output graph:
    input(shape=10,2,3)----->relu---->out1
                        |
                        |
                        --->reduce(axis=1)----->log---->out2
    """

    def test_fusion_with_axis_op(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 2, 3))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 2, 1])
            x1 = mb.relu(x=x)
            x2 = mb.reduce_mean(x=x, axes=[2], keep_dims=True)
            y1 = mb.transpose(x=x1, perm=[0, 2, 1])
            x3 = mb.log(x=x2)
            y2 = mb.transpose(x=x3, perm=[0, 2, 1])
            return y1, y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "reduce_mean", "transpose", "log", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "reduce_mean", "log"])

        assert list(block.find_ops(op_type="reduce_mean")[0].inputs["axes"].val) == [1]
        assert_model_is_valid(
            prog,
            {"x": (10, 2, 3)},
            expected_output_shapes={
                block.outputs[0].name: (10, 2, 3),
                block.outputs[1].name: (10, 1, 3),
            },
        )

    """
    Input graph:
    input(shape=11,2,3,6)--->transpose(axis=[0,3,1,2])---
                                                       |
                                                       |
                                                        --->pad(pad=[0,0,0,0,1,2,3,4])
                                                              |
                                                              |-->log--->transpose(axis=[0,2,3,1])-->out1(shape=11,5,10,6)

    Output graph:
    same as input graph, as transpose cannot be pushed through the pad op since "reflect" mode is only supported
    along the last two axis
    """

    def test_fusion_with_pad_reflective_op_0(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(11, 2, 3, 6))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x2 = mb.pad(x=x, pad=[0, 0, 0, 0, 1, 2, 3, 4], mode="reflect")
            x3 = mb.log(x=x2)
            y2 = mb.transpose(x=x3, perm=[0, 2, 3, 1])
            return y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog), ["transpose", "pad", "log", "transpose"]
        )
        self.assertEqual(get_op_types_in_program(prog), ["transpose", "pad", "log", "transpose"])

        assert list(block.find_ops(op_type="pad")[0].inputs["pad"].val.flatten()) == [
            0,
            0,
            0,
            0,
            1,
            2,
            3,
            4,
        ]
        assert_model_is_valid(
            prog,
            {"x": (11, 2, 3, 6)},
            expected_output_shapes={block.outputs[0].name: (11, 5, 10, 6)},
        )

    """
    Input graph:
    input(shape=11,2,3,6)--->transpose(axis=[0,1,3,2])---
                                                       |
                                                       |
                                                        --->pad(pad=[0,0,0,0,1,2,3,4])
                                                              |
                                                              |-->log--->transpose(axis=[0,1,3,2])-->out1(shape=11,2,10,9)

    Output graph:
    input(shape=11,2,3,6)--->pad(pad=[0,0,0,0,3,4,1,2])-->log-->out1(shape=11,2,10,9)
    """

    def test_fusion_with_pad_reflective_op_1(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(11, 2, 3, 6))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 1, 3, 2])
            x2 = mb.pad(x=x, pad=[0, 0, 0, 0, 1, 2, 3, 4], mode="reflect")
            x3 = mb.log(x=x2)
            y2 = mb.transpose(x=x3, perm=[0, 1, 3, 2])
            return y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog), ["transpose", "pad", "log", "transpose"]
        )
        self.assertEqual(get_op_types_in_program(prog), ["pad", "log"])

        assert list(block.find_ops(op_type="pad")[0].inputs["pad"].val.flatten()) == [
            0,
            0,
            0,
            0,
            3,
            4,
            1,
            2,
        ]
        assert_model_is_valid(
            prog,
            {"x": (11, 2, 3, 6)},
            expected_output_shapes={block.outputs[0].name: (11, 2, 10, 9)},
        )

    """
    Input graph:
    input(shape=11,2,3,6)--->transpose(axis=[0,3,1,2])---
                                                       |
                                                       |
                                                        --->pad(pad=[0,0,0,0,1,2,3,4])
                                                              |
                                                              |-->log--->transpose(axis=[0,2,3,1])-->out1(shape=11,5,10,6)

    Output graph:
    input(shape=11,2,3,6)--->pad(pad=[0,0,1,2,3,4,0,0])-->log-->out1(shape=11,5,10,6)
    """

    def test_fusion_with_pad_constant_op(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(11, 2, 3, 6))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x2 = mb.pad(x=x, pad=[0, 0, 0, 0, 1, 2, 3, 4], mode="constant", constant_val=3.0)
            x3 = mb.log(x=x2)
            y2 = mb.transpose(x=x3, perm=[0, 2, 3, 1])
            return y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog), ["transpose", "pad", "log", "transpose"]
        )
        self.assertEqual(get_op_types_in_program(prog), ["pad", "log"])

        assert list(block.find_ops(op_type="pad")[0].inputs["pad"].val.flatten()) == [
            0,
            0,
            1,
            2,
            3,
            4,
            0,
            0,
        ]
        assert_model_is_valid(
            prog,
            {"x": (11, 2, 3, 6)},
            expected_output_shapes={block.outputs[0].name: (11, 5, 10, 6)},
        )

    """
    Input graph:
                                                    const(shape=2)
                                                          |
                                                          V
    input(shape=1,2,5,5)--->transpose(axis=[0,2,3,1])--->add---->transpose(axis=[0,3,1,2])--->out(shape=1,2,5,5)

    Output graph:
                        const(shape=1,2,1,1)
                             |
                             V
    input(shape=1,2,5,5)--->add--->out(shape=1,2,5,5)
    """

    def test_fusion_with_add_constant_op(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.add(x=x, y=np.array([10.0, 100.0]))
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(get_op_types_in_program(prev_prog), ["transpose", "add", "transpose"])
        self.assertEqual(get_op_types_in_program(prog), ["add"])

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 5, 5)},
        )

    """
    Input graph:
                                                    const(scalar)
                                                          |
                                                          V
    input(shape=1,2,5,5)--->transpose(axis=[0,2,3,1])--->add---->transpose(axis=[0,3,1,2])--->out(shape=1,2,5,5)

    Output graph:
                        const(scalar)
                             |
                             V
    input(shape=1,2,5,5)--->add--->out(shape=1,2,5,5)
    """

    def test_fusion_with_add_scalar_constant_op(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.add(x=5.0, y=x)
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(get_op_types_in_program(prev_prog), ["transpose", "add", "transpose"])
        self.assertEqual(get_op_types_in_program(prog), ["add"])

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 5, 5)},
        )

    """
    Input graph:
    input(shape=1,2,5,5)----->transpose(axis=[0,2,3,1])--->add---->transpose(axis=[0,3,1,2])--->out(shape=1,2,5,5)
                          |                                 ^
                          |                                 |
                          |---->relu---->transpose(axis=[0,2,3,1])

    Output graph:
    input(shape=1,2,5,5)----->add--->out(shape=1,2,5,5)
                      |        ^
                      |        |
                      |------>relu
    """

    def test_fusion_with_add_broadcastable_0(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.relu(x=x)
            x2 = mb.transpose(x=x2, perm=[0, 2, 3, 1])
            x3 = mb.add(x=x1, y=x2)
            y = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            return y

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "transpose", "add", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "add"])

        assert block.find_ops(op_type="relu")[0].inputs["x"] == block.inputs["x"]
        assert block.find_ops(op_type="add")[0].inputs["x"] == block.inputs["x"]
        assert (
            block.find_ops(op_type="add")[0].inputs["y"]
            == block.find_ops(op_type="relu")[0].outputs[0]
        )

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 5, 5)},
        )

    """
    Input graph:
    input(shape=1,2,5,5)----->transpose(axis=[0,2,3,1])--->add---->transpose(axis=[0,3,1,2])--->out(shape=1,2,5,5)
                          |                                 ^
                          |                                 |
                          |----------------------->transpose(axis=[0,2,3,1])

    Output graph:
    input(shape=1,2,5,5)----->add--->out(shape=1,2,5,5)
                      |        ^
                      |        |
                      |---------
    """

    def test_fusion_with_add_broadcastable_1(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x3 = mb.add(x=x1, y=x2)
            y = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            return y

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "transpose", "add", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["add"])

        assert block.find_ops(op_type="add")[0].inputs["x"] == block.inputs["x"]
        assert block.find_ops(op_type="add")[0].inputs["y"] == block.inputs["x"]

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 5, 5)},
        )

    """
    Input graph:
    input(shape=1,2,5,5)--->transpose(axis=[0,2,3,1])---> relu---->concat(axis=3)----->transpose(axis=[0,3,1,2])----->out1(shape=1,4,5,5)
                         |                                              ^
                         |                                              |
                         |->transpose(axis=[0,2,3,1])--->relu------------

    Output graph:
    input(shape=1,2,5,5)------> relu---->concat(axis=1)--->out1(shape=1,4,5,5)
                         |                    ^
                         |                    |
                         |---->relu------------
    """

    def test_concat_pattern_0(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x1 = mb.relu(x=x1)
            x2 = mb.relu(x=x2)
            x3 = mb.concat(values=[x1, x2], axis=3)
            x4 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            return x4

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "transpose", "relu", "relu", "concat", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "relu", "concat"])

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={block.outputs[0].name: (1, 4, 5, 5)},
        )

    """
    Input graph:
    input(shape=1,2,5,5)--->transpose(axis=[0,2,3,1])---> relu---->concat(axis=3)----->transpose(axis=[0,3,1,2])----->out1(shape=1,4,5,5)
                         |                                              ^
                         |                                              |
                         |->transpose(axis=[0,2,3,1])------->relu--------
                                                        |
                                                        V
                                                       pool--->out2(shape=1,5,5,2)



    Output graph:
    input(shape=1,2,5,5)------> relu---->concat(axis=1)--->out1(shape=1,4,5,5)
                         |                    ^
                         |                    |
                         |---->relu------------
                         |
                         |--->transpose(axis=[0,2,3,1])---->pool--->out2(shape=1,5,5,2)
    """

    def test_concat_pattern_1(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x1 = mb.relu(x=x1)
            x2 = mb.relu(x=x2)
            x3 = mb.concat(values=[x1, x2], axis=3)
            x4 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            x5 = mb.avg_pool(x=x2, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            return x4, x5

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "transpose",
                "transpose",
                "relu",
                "relu",
                "concat",
                "transpose",
                "avg_pool",
            ],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["relu", "relu", "concat", "transpose", "avg_pool"],
        )
        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={
                block.outputs[0].name: (1, 4, 5, 5),
                block.outputs[1].name: (1, 5, 5, 2),
            },
        )

    """
    Input graph:
    input(shape=1,2,5,5)--->transpose(axis=[0,2,3,1])---> relu---->concat(axis=3)----->transpose(axis=[0,3,1,2])----->out1(shape=1,4,5,5)
                         |                                              ^
                         |                                              |
                         |->transpose(axis=[0,2,3,1])------->relu--------
                                                        |
                                                        V
                                                       relu--->out2(shape=1,5,5,2)



    Output graph:
    input(shape=1,2,5,5)------> relu---->concat(axis=1)--->out1(shape=1,4,5,5)
                         |                    ^
                         |                    |
                         |---->relu------------
                         |
                         |--->relu---->transpose(axis=[0,2,3,1])---->out2(shape=1,5,5,2)
    """

    def test_concat_pattern_2(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x1 = mb.relu(x=x1)
            x2 = mb.relu(x=x2)
            x3 = mb.concat(values=[x1, x2], axis=3)
            x4 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            x5 = mb.relu(x=x2)
            return x4, x5

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "transpose", "relu", "relu", "concat", "transpose", "relu"],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["relu", "relu", "concat", "relu", "transpose"],
        )

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={
                block.outputs[0].name: (1, 4, 5, 5),
                block.outputs[1].name: (1, 5, 5, 2),
            },
        )

    """
    Input graph:
    input(shape=1,2,5,5)--->transpose(axis=[0,2,3,1])---> relu---->concat(axis=3)----->transpose(axis=[0,3,1,2])----->out1(shape=1,4,5,5)
                         |                                              ^
                         |                                              |
                         |->transpose(axis=[0,2,3,1])------->relu--------
                                                        |
                                                        V
                                                       out2(shape=1,5,5,2)



    Output graph:
    input(shape=1,2,5,5)------> relu---->concat(axis=1)--->out1(shape=1,4,5,5)
                         |                    ^
                         |                    |
                         |---->relu------------
                         |
                         |--->transpose(axis=[0,2,3,1])---->out2(shape=1,5,5,2)
    """

    def test_concat_pattern_3(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x1 = mb.relu(x=x1)
            x2 = mb.relu(x=x2)
            x3 = mb.concat(values=[x1, x2], axis=3)
            x4 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            return x4, x2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "transpose", "relu", "relu", "concat", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "relu", "concat", "transpose"])

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={
                block.outputs[0].name: (1, 4, 5, 5),
                block.outputs[1].name: (1, 5, 5, 2),
            },
        )

    """
    Input graph:
    input(shape=1,2,5,5)--->transpose(axis=[0,2,3,1])---> relu---->concat(axis=3)----->transpose(axis=[0,3,1,2])----->out1(shape=1,4,5,5)
                         |                                              ^
                         |                                              |
                         |->transpose(axis=[0,2,3,1])------->relu--------
                                                        |
                                                        V
                                            transpose(axis=[0,3,1,2]) -----> out2(shape=1,2,5,5)

    Output graph:
    input(shape=1,2,5,5)---> relu---->concat(axis=1)----->out1(shape=1,4,5,5)
                     |                     ^
                     |                     |
                     |------------------->relu-------->out2(shape=1,2,5,5)
    """

    def test_concat_pattern_4(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x1 = mb.relu(x=x1)
            x2 = mb.relu(x=x2)
            x3 = mb.concat(values=[x1, x2], axis=3)
            x4 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            x5 = mb.transpose(x=x2, perm=[0, 3, 1, 2])
            return x4, x5

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "transpose",
                "transpose",
                "relu",
                "relu",
                "concat",
                "transpose",
                "transpose",
            ],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "relu", "concat"])

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={
                block.outputs[0].name: (1, 4, 5, 5),
                block.outputs[1].name: (1, 2, 5, 5),
            },
        )

    """
    Input graph:
                                                     constant(shape=[30,10,5])
                                                            |
                                                            V
    input(shape=10,20,30)--->transpose(axis=[2,0,1])--->concat(axis=2)----->transpose(axis=[1,2,0])----->out1(shape=10,25,30)

    Output graph:
                        constant(shape=[10,5,30])
                                |
                                V
    input(shape=10,20,30)--->concat(axis=1)----->out1(shape=10,25,30)
    """

    def test_concat_pattern_5(self):
        const = np.random.rand(30, 10, 5)

        @mb.program(input_specs=[mb.TensorSpec(shape=(10, 20, 30))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[2, 0, 1])
            c = mb.const(val=const)
            x2 = mb.concat(values=[x1, c], axis=2)
            x3 = mb.transpose(x=x2, perm=[1, 2, 0])
            return x3

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(get_op_types_in_program(prev_prog), ["transpose", "concat", "transpose"])
        self.assertEqual(get_op_types_in_program(prog), ["concat"])

        assert_model_is_valid(
            prog,
            {"x": (10, 20, 30)},
            expected_output_shapes={block.outputs[0].name: (10, 25, 30)},
        )

    """
    Input graph:
                                        input2(shape=30,10,20)-----|
                                                                   |
    input(shape=10,20,30)--->transpose(axis=[2,0,1])----->relu-----|----->concat(axis=2)------>out1(shape=90,10,20)
                                                      |            |
                                                      |-->relu-----|
                                                      |
                                                      |-->relu---->transpose(axis=[1,2,0])---->out2(shape=10,20,30)
                                                      |
                                                      |-->relu---->transpose(axis=[1,2,0])---->out3(shape=10,20,30)
                                                      |
                                                      |-->relu---->transpose(axis=[1,2,0])---->out4(shape=10,20,30)

    Output graph:

                                        input2(shape=30,10,20)-----|
                                                                   |
    input(shape=10,20,30)----->relu--->transpose(axis=[2,0,1])-----|----->concat(axis=2)------>out1(shape=90,10,20)
                           |                                       |
                           |-->relu--->transpose(axis=[2,0,1])-----|
                           |
                           |-->relu---->out2(shape=10,20,30)
                           |
                           |-->relu---->out3(shape=10,20,30)
                           |
                           |-->relu---->out4(shape=10,20,30)

    Output graph:
    """

    def test_concat_pattern_6(self):
        @mb.program(
            input_specs=[
                mb.TensorSpec(shape=(10, 20, 30)),
                mb.TensorSpec(shape=(30, 10, 20)),
            ]
        )
        def prog(x, y):
            x1 = mb.transpose(x=x, perm=[2, 0, 1])
            r1 = mb.relu(x=x1)
            r2 = mb.relu(x=x1)
            r3 = mb.relu(x=x1)
            r4 = mb.relu(x=x1)
            r5 = mb.relu(x=x1)

            x2 = mb.concat(values=[r1, r2, y], axis=0)
            x3 = mb.transpose(x=r3, perm=[1, 2, 0])
            x4 = mb.transpose(x=r4, perm=[1, 2, 0])
            x5 = mb.transpose(x=r5, perm=[1, 2, 0])
            return x2, x3, x4, x5

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "transpose",
                "relu",
                "relu",
                "relu",
                "relu",
                "relu",
                "concat",
                "transpose",
                "transpose",
                "transpose",
            ],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            [
                "relu",
                "relu",
                "relu",
                "relu",
                "relu",
                "transpose",
                "transpose",
                "concat",
            ],
        )

        assert_model_is_valid(
            prog,
            {"x": (10, 20, 30), "y": (30, 10, 20)},
            expected_output_shapes={
                block.outputs[0].name: (90, 10, 20),
                block.outputs[1].name: (10, 20, 30),
                block.outputs[2].name: (10, 20, 30),
                block.outputs[3].name: (10, 20, 30),
            },
        )

    """
    Input graph:
    input(shape=1,4,5,6)--->transpose(axis=[0,3,2,1])--->relu---->split(axis=1, num_splits=2)----->transpose(axis=[0,3,2,1])----->out1(shape=1,4,5,3)
                                                                        |
                                                                        v
                                                                    transpose(axis[0,3,2,1])-------------------------->out2(shape=1,4,5,3)

    Output graph:
    input(shape=1,4,5,6)------> relu ---->split(axis=3)--->out1(shape=1,4,5,3)
                                              |
                                              v
                                            out2(shape=1,4,5,3)
    """

    def test_split_nd_pattern_0(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 4, 5, 6))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 3, 2, 1])
            x1 = mb.relu(x=x1)
            x2, x3 = mb.split(x=x1, axis=1, num_splits=2)
            x4 = mb.transpose(x=x2, perm=[0, 3, 2, 1])
            x5 = mb.transpose(x=x3, perm=[0, 3, 2, 1])
            return x4, x5

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "split", "transpose", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "split"])

        assert_model_is_valid(
            prog,
            {"x": (1, 4, 5, 6)},
            expected_output_shapes={
                block.outputs[0].name: (1, 4, 5, 3),
                block.outputs[1].name: (1, 4, 5, 3),
            },
        )

        self.assertEqual(block.find_ops(op_type="split")[0].axis.val, 3)

    """
    Input graph:
    input(shape=1,4,5,6)--->transpose(axis=[0,3,2,1])--->relu---->splitd(axis=1, num_splits=6)----->transpose(axis=[0,3,2,1])----->out1(shape=1,4,5,3)
                                                                        |
                                                                        v
                                                                    transpose(axis[0,3,2,1])-------------------------------------->out2(shape=1,4,5,3)

    Output graph:
    input(shape=1,4,5,6)------>relu---->split(axis=3)--->out1(shape=1,4,5,3)
                                              |
                                              v
                                            out2(shape=1,4,5,3)
    """

    def test_split_nd_pattern_1(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 4, 5, 6))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 3, 2, 1])
            x1 = mb.relu(x=x1)
            x2, x3, x4, x5, x6, x7 = mb.split(x=x1, axis=1, num_splits=6)
            x2 = mb.transpose(x=x2, perm=[0, 3, 2, 1])
            x3 = mb.transpose(x=x3, perm=[0, 3, 2, 1])
            x4 = mb.transpose(x=x4, perm=[0, 3, 2, 1])
            x5 = mb.transpose(x=x5, perm=[0, 3, 2, 1])
            x6 = mb.transpose(x=x6, perm=[0, 3, 2, 1])
            x7 = mb.transpose(x=x7, perm=[0, 3, 2, 1])
            return x2, x3, x4, x5, x6, x7

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "transpose",
                "relu",
                "split",
                "transpose",
                "transpose",
                "transpose",
                "transpose",
                "transpose",
                "transpose",
            ],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "split"])

        assert_model_is_valid(
            prog,
            {"x": (1, 4, 5, 6)},
            expected_output_shapes={
                block.outputs[0].name: (1, 4, 5, 1),
                block.outputs[1].name: (1, 4, 5, 1),
                block.outputs[2].name: (1, 4, 5, 1),
                block.outputs[3].name: (1, 4, 5, 1),
                block.outputs[4].name: (1, 4, 5, 1),
                block.outputs[5].name: (1, 4, 5, 1),
            },
        )

        self.assertEqual(block.find_ops(op_type="split")[0].axis.val, 3)

    """
    Input graph:
    input(shape=1,4,5,6)--->transpose(axis=[0,3,2,1])---> split(axis=1, num_splits=2) ----> concat(axis=1) ----->transpose(axis=[0,3,2,1]) ----->out1(shape=1,4,5,6)
                                                                        |                         ^
                                                                        v                         |
                                                                      relu() ----------------------

    Output graph:
    input(shape=1,4,5,6)------>split(axis=3)--->concat(axis=3) -------> out1(shape=1,4,5,6)
                                        |               ^
                                        v               |
                                    relu() --------------
    """

    def test_split_nd_pattern_2(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 4, 5, 6))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 3, 2, 1])
            x2, x3 = mb.split(x=x1, axis=1, num_splits=2)
            x4 = mb.relu(x=x2)
            x5 = mb.concat(values=[x4, x3], axis=1)
            x6 = mb.transpose(x=x5, perm=[0, 3, 2, 1])
            return x6

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "split", "relu", "concat", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["split", "relu", "concat"])

        assert_model_is_valid(
            prog,
            {"x": (1, 4, 5, 6)},
            expected_output_shapes={block.outputs[0].name: (1, 4, 5, 6)},
        )

        self.assertEqual(block.find_ops(op_type="split")[0].axis.val, 3)

    """
    Input graph:
    input(shape=1,5,5,3)----->transpose(axis=[0,3,1,2])
                                |
                                ---->relu-------------->transpose(axis=[0,2,3,1])
                                             |                  |
                                             |                  V
                                             |                 relu
                                             |                  |
                                             |                  V
                                             |         transpose(axis=[0,3,1,2])
                                             |                  |
                                             |                  V
                                             ----------------> add --------> relu---->pool---->out(shape=1,3,5,5)


    Output graph:


    input(shape=1,5,5,3)---->relu------------------------> relu
                                         |                  |
                                         |                  V
                                         ----------------> add
                                                            |
                                                            V
                                                           relu
                                                            |
                                                            V
                                                           transpose(axis=[0,3,1,2])-->pool---->out(shape=1,3,5,5)

    """

    def test_skip_connection_pattern_0(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 5, 5, 3))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.relu(x=x)
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.relu(x=x1)
            x3 = mb.transpose(x=x2, perm=[0, 3, 1, 2])
            x4 = mb.add(x=x, y=x3)
            x5 = mb.relu(x=x4)
            x6 = mb.avg_pool(x=x5, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            return x6

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "transpose",
                "relu",
                "transpose",
                "relu",
                "transpose",
                "add",
                "relu",
                "avg_pool",
            ],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["relu", "relu", "add", "relu", "transpose", "avg_pool"],
        )
        assert_model_is_valid(
            prog,
            {"x": (1, 5, 5, 3)},
            expected_output_shapes={block.outputs[0].name: (1, 3, 5, 5)},
        )

    """
    Input graph:
    input(shape=1,5,5,3)----->transpose(axis=[0,3,1,2])
                                |
                                ---->relu-------------->transpose(axis=[0,2,3,1])
                                             |                  |
                                             |                  V
                                             |                 relu
                                             |                  |
                                             |                  V
                                             |         transpose(axis=[0,3,1,2])
                                             |                  |
                                             |                  V
                                             ----------------> add -->transpose(axis=[0,2,3,1])
                                                                            |
                                                                            V
                                                                          relu---->pool---->out(shape=1,5,5,3)


    Output graph:


    input(shape=1,5,5,3)---->relu------------------------> relu
                                         |                  |
                                         |                  V
                                         ----------------> add
                                                            |
                                                            V
                                                           relu
                                                            |
                                                            V
                                                           pool---->out(shape=1,5,5,3)

    """

    def test_skip_connection_pattern_1(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 5, 5, 3))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.relu(x=x)
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.relu(x=x1)
            x3 = mb.transpose(x=x2, perm=[0, 3, 1, 2])
            x4 = mb.add(x=x, y=x3)
            x4 = mb.transpose(x=x4, perm=[0, 2, 3, 1])
            x5 = mb.relu(x=x4)
            x6 = mb.avg_pool(x=x5, kernel_sizes=[1, 1], strides=[1, 1], pad_type="valid")
            return x6

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "transpose",
                "relu",
                "transpose",
                "relu",
                "transpose",
                "add",
                "transpose",
                "relu",
                "avg_pool",
            ],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "relu", "add", "relu", "avg_pool"])
        assert_model_is_valid(
            prog,
            {"x": (1, 5, 5, 3)},
            expected_output_shapes={block.outputs[0].name: (1, 5, 5, 3)},
        )

    """
    Input graph:
    input(shape=2,5)--->transpose(axis=[1,0])--->transpose(axis=[1,0])-->reduce(axis=1)
                                |                                             |
                                |                                             V
                                |                                    transpose(axis=[1,0])
                                |                                             |
                                |                                             V
                                -------------------------------------------->add------->out(shape=5,2)

    Output graph:
    input(shape=2,5)--->reduce(axis=1)---->add---->transpose(axis=[1,0])--->out(shape=5,2)
                     |                      ^
                     |                      |
                     ------------------------
    """

    def test_residual_with_unmaterialized_output(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[1, 0])
            t1 = mb.transpose(x=x1, perm=[1, 0])
            x2 = mb.reduce_mean(x=t1, axes=[1], keep_dims=True)
            t2 = mb.transpose(x=x2, perm=[1, 0])
            return mb.add(x=x1, y=t2)

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "transpose", "reduce_mean", "transpose", "add"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["reduce_mean", "add", "transpose"])

        assert_model_is_valid(
            prog, {"x": (2, 5)}, expected_output_shapes={block.outputs[0].name: (5, 2)}
        )

    """
    Input graph:
    input(shape=2,5)--->transpose(axis=[1,0])--->transpose(axis=[1,0])-->reduce(axis=1)
                                |                                             |
                                |                                             V
                                |                                    transpose(axis=[1,0])
                                |                                             |
                                |                                             V
                                -------------------------------------------->add------->out1(shape=5,2)
                                                                              |
                                                                              V
                                                                            relu------->out2(shape=5,2)

    Output graph:
    input(shape=2,5)--->reduce(axis=1)----> add ----->transpose(axis=[1,0])----->out1(shape=5,2)
                     |                       |
                     |                       V
                     ---------------------> relu----->transpose(axis=[1,0])----->out2(shape=5,2)
    """

    def test_residual_with_unmaterialized_multiple_output(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[1, 0])
            t1 = mb.transpose(x=x1, perm=[1, 0])
            x2 = mb.reduce_mean(x=t1, axes=[1], keep_dims=True)
            t2 = mb.transpose(x=x2, perm=[1, 0])
            out1 = mb.add(x=x1, y=t2)
            out2 = mb.relu(x=out1)
            return out1, out2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "transpose", "reduce_mean", "transpose", "add", "relu"],
        )
        self.assertEqual(
            get_op_types_in_program(prog), ["reduce_mean", "add", "relu", "transpose", "transpose"]
        )

        assert_model_is_valid(
            prog,
            {"x": (2, 5)},
            expected_output_shapes={block.outputs[0].name: (5, 2), block.outputs[1].name: (5, 2)},
        )

    """
    Input graph:
    input(shape=2,5)---->transpose(axis=[1,0])------>relu----->transpose(axis=[1,0])------>out2(shape=2,5)
                                                       |
                                                       ------->out1(shape=5,2)

    Output graph:
    input(shape=2,5)---->relu-----> out2(shape=2,5)
                           |
                           V
                    transpose(axis=[1,0]) -----> out1(shape=5,2)
    """

    def test_materialized_output_reuse(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[1, 0])
            y1 = mb.relu(x=x1)
            y2 = mb.transpose(x=y1, perm=[1, 0])
            return y1, y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            [
                "transpose",
                "relu",
                "transpose",
            ],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "transpose"])

        assert_model_is_valid(
            prog,
            {"x": (2, 5)},
            expected_output_shapes={block.outputs[0].name: (5, 2), block.outputs[1].name: (2, 5)},
        )

    """
    Input graph:
    input(shape=1,2,5,5)----->transpose(axis=[0,2,3,1])------->add------------>transpose(axis=[0,3,1,2])--->out1(shape=1,2,5,5)
                                                          |     ^        |
                                                          |     |        |
                                                          ---->relu      ----->transpose(axis=[0,3,1,2])--->out2(shape=1,2,5,5)

    Output graph:
    input(shape=1,2,5,5)----->add------->out1(shape=1,2,5,5)
                      |        ^    |
                      |        |    |
                      |------>relu  ------identity(renaming)---->out2(shape=1,2,5,5)
    """

    def test_fusion_with_double_outputs(self):
        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 2, 5, 5))])
        def prog(x):
            x1 = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x2 = mb.relu(x=x1)
            x3 = mb.add(x=x1, y=x2)
            y1 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            y2 = mb.transpose(x=x3, perm=[0, 3, 1, 2])
            return y1, y2

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "relu", "add", "transpose", "transpose"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["relu", "add", "identity"])

        assert block.find_ops(op_type="relu")[0].inputs["x"] == block.inputs["x"]
        assert block.find_ops(op_type="add")[0].inputs["x"] == block.inputs["x"]
        assert (
            block.find_ops(op_type="add")[0].inputs["y"]
            == block.find_ops(op_type="relu")[0].outputs[0]
        )

        assert_model_is_valid(
            prog,
            {"x": (1, 2, 5, 5)},
            expected_output_shapes={block.outputs[0].name: (1, 2, 5, 5)},
        )

    def test_pass_through_broadcasted_binary_op(self):
        """
        Input graph:
                                                                const (shape=(1,1,1,3))
                                                                     |
        input (shape=(1,4,3,2)) --> transpose (shape=(1,2,4,3)) --> add --> transpose --> relu

        Output graph:

                             const (shape=(1,1,3,1))
                                     |
        input (shape=(1,4,3,2)) --> add --> relu
        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 4, 3, 2))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.add(x=x, y=np.array(np.ones(shape=(1, 1, 1, 3))))
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "add", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["add", "relu"])
        assert_model_is_valid(
            prog,
            {"x": (1, 4, 3, 2)},
            expected_output_shapes={block.outputs[0].name: (1, 4, 3, 2)},
        )

    def test_binary_op_with_constant_input(self):
        """
        Input graph:
                                                                const (shape=(4,3))
                                                                     |
        input (shape=(1,4,3,2)) --> transpose (shape=(1,2,4,3)) --> add --> transpose --> relu

        Output graph:

                             const (shape=(1,4,3,1))
                                     |
        input (shape=(1,4,3,2)) --> add --> relu
        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 4, 3, 2))])
        def prog(x):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.add(x=x, y=np.array(np.ones(shape=(4, 3))))
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "add", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["add", "relu"])
        assert_model_is_valid(
            prog,
            {"x": (1, 4, 3, 2)},
            expected_output_shapes={block.outputs[0].name: (1, 4, 3, 2)},
        )

    def test_binary_op_with_non_constant_input1(self):
        """
        Input graph:
                                                                input (shape=(3,))
                                                                     |
        input (shape=(1,4,3,2)) --> transpose (shape=(1,2,4,3)) --> add --> transpose --> relu

        Output graph:

                                  input (shape=(3,))
                                     |
                                 reshape (shape=(1,1,3,1))
                                     |
        input (shape=(1,4,3,2)) --> add --> relu
        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(1, 4, 3, 2)), mb.TensorSpec(shape=(3,))])
        def prog(x, y):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.add(x=x, y=y)
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "add", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["reshape", "add", "relu"])
        reshape_op = prog.find_ops(op_type="reshape", exactly_one=True)[0]
        assert reshape_op.outputs[0].shape == (1, 1, 3, 1)
        assert_model_is_valid(
            prog,
            {"x": (1, 4, 3, 2), "y": (3,)},
            expected_output_shapes={block.outputs[0].name: (1, 4, 3, 2)},
        )

    def test_binary_op_with_non_constant_input2(self):
        """
        Input graph:
                                                                input (shape=(3,1,2))
                                                                     |
        input (shape=(5,3,4,2)) --> transpose (shape=(4,3,5,2)) --> add --> transpose --> relu

        Output graph:

                                  input (shape=(3,1,2))
                                     |
                                 reshape (shape=(1,3,1,2))
                                     |
        input (shape=(5,3,4,2)) --> add --> relu
        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(5, 3, 4, 2)), mb.TensorSpec(shape=(3, 1, 2))])
        def prog(x, y):
            x = mb.transpose(x=x, perm=[2, 1, 0, 3])
            x = mb.add(x=x, y=y)
            x = mb.transpose(x=x, perm=[2, 1, 0, 3])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "add", "transpose", "relu"],
        )
        self.assertEqual(get_op_types_in_program(prog), ["reshape", "add", "relu"])
        reshape_op = prog.find_ops(op_type="reshape", exactly_one=True)[0]
        assert reshape_op.outputs[0].shape == (1, 3, 1, 2)
        assert_model_is_valid(
            prog,
            {"x": (5, 3, 4, 2), "y": (3, 1, 2)},
            expected_output_shapes={block.outputs[0].name: (5, 3, 4, 2)},
        )

    def test_binary_op_with_non_constant_input3(self):
        """
        Input graph:
                                                                input (shape=(3,1,2))
                                                                     |
        input (shape=(s,3,4,2)) --> transpose (shape=(4,3,s,2)) --> add --> transpose --> relu

        Output graph:

                                input (shape=(3,1,2))
                                     |
                                 reshape (shape=(1,3,1,2))
                                     |
        input (shape=(s,3,4,2)) --> add --> relu
        """

        @mb.program(
            input_specs=[
                mb.TensorSpec(shape=(get_new_symbol(), 3, 4, 2)),
                mb.TensorSpec(shape=(3, 1, 2)),
            ]
        )
        def prog(x, y):
            x = mb.transpose(x=x, perm=[2, 1, 0, 3])
            x = mb.add(x=x, y=y)
            x = mb.transpose(x=x, perm=[2, 1, 0, 3])
            x = mb.relu(x=x)
            return x

        pass_name = "common::reduce_transposes"
        PASS_REGISTRY[pass_name](prog)
        self.assertEqual(get_op_types_in_program(prog), ["reshape", "add", "relu"])
        reshape_op = prog.find_ops(op_type="reshape", exactly_one=True)[0]
        assert reshape_op.outputs[0].shape == (1, 3, 1, 2)
        block = prog.functions["main"]
        assert_model_is_valid(
            prog,
            {"x": (5, 3, 4, 2), "y": (3, 1, 2)},
            expected_output_shapes={block.outputs[0].name: (5, 3, 4, 2)},
        )

    def test_binary_op_with_non_constant_input4(self):
        """
        Input graph:
                                                                input (shape=(3,s,2))
                                                                     |
        input (shape=(1,3,4,2)) --> transpose (shape=(4,3,1,2)) --> add --> transpose --> relu

        Output graph: same as input graph since the non-transpose input of the add op has symbolic shape
        """

        @mb.program(
            input_specs=[
                mb.TensorSpec(shape=(1, 3, 4, 2)),
                mb.TensorSpec(shape=(3, get_new_symbol(), 2)),
            ]
        )
        def prog(x, y):
            x = mb.transpose(x=x, perm=[2, 1, 0, 3])
            x = mb.add(x=x, y=y)
            x = mb.transpose(x=x, perm=[2, 1, 0, 3])
            x = mb.relu(x=x)
            return x

        pass_name = "common::reduce_transposes"
        PASS_REGISTRY[pass_name](prog)
        self.assertEqual(get_op_types_in_program(prog), ["transpose", "add", "transpose", "relu"])
        block = prog.functions["main"]
        assert_model_is_valid(
            prog,
            {"x": (1, 3, 4, 2), "y": (3, 10, 2)},
            expected_output_shapes={block.outputs[0].name: (10, 3, 4, 2)},
        )

    def test_binary_op_with_non_constant_input5(self):
        """
        Input graph:
                                                                input (shape=(3,4))
                                                                     |
        input (shape=(5,3,4,2)) --> transpose (shape=(5,2,3,4)) --> add --> transpose --> relu

        Output graph: same as input graph since transpose compliment for 2nd input of add cannot be represented
        as a static reshape
        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(5, 3, 4, 2)), mb.TensorSpec(shape=(3, 4))])
        def prog(x, y):
            x = mb.transpose(x=x, perm=[0, 3, 1, 2])
            x = mb.add(x=x, y=y)
            x = mb.transpose(x=x, perm=[0, 2, 3, 1])
            x = mb.relu(x=x)
            return x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")
        self.assertEqual(
            get_op_types_in_program(prev_prog),
            ["transpose", "add", "transpose", "relu"],
        )
        self.assertEqual(
            get_op_types_in_program(prog),
            ["transpose", "add", "transpose", "relu"],
        )
        assert_model_is_valid(
            prog,
            {"x": (5, 3, 4, 2), "y": (3, 4)},
            expected_output_shapes={block.outputs[0].name: (5, 3, 4, 2)},
        )

    def test_input_duplicate_output(self):
        """
        Input graph:
        input -----> out (consist of duplicated input)

        Output graph:
        input -----> out (consist of duplicated input)

        Notice that a temp identity sink is added for all outputs, so the block before going through the pass is:
            function[CoreML3](%x: (2, 2, 1, 1, fp32)(Tensor)) {
              block0() {
                %identity_0: (2, 2, 1, 1, fp32)(Tensor) = identity(x=%x, name="identity_0")
              } -> (%identity_0, %identity_0)
            }
        """

        @mb.program(input_specs=[mb.TensorSpec(shape=(2, 2, 1, 1))])
        def prog(x):
            return x, x

        prev_prog, prev_block, block = apply_pass_and_basic_check(prog, "common::reduce_transposes")

        self.assertEqual(get_op_types_in_program(prev_prog), [])
        self.assertEqual(get_op_types_in_program(prog), [])
        assert_model_is_valid(
            prog,
            {"x": (2, 2, 1, 1)},
            backend=("mlprogram", "fp16"),
            expected_output_shapes={
                block.outputs[0].name: (2, 2, 1, 1),
                block.outputs[1].name: (2, 2, 1, 1),
            },
        )


class TestTransposePassUtilityMethods:
    @staticmethod
    @pytest.mark.parametrize("rank", [1, 2, 3, 4, 5])
    def test_transpose_compliment_method(rank):
        x = np.random.rand(*np.random.randint(low=1, high=15, size=rank))
        perm = np.random.permutation(rank)
        reverse_perm = TransformAxisUpdateOps._find_transpose_compliment(perm)
        x_transpose = np.transpose(x, perm)
        x_transpose_transpose = np.transpose(x_transpose, reverse_perm)
        np.testing.assert_equal(x, x_transpose_transpose)
