code: def test_simplify_transpose():
      # Test a series of transpose and layout_transform ops
      def before1():
          x = relay.var(
              "x", shape=(1, 3, 224, 224), dtype="float32"
          )  # NCHW
          y = relay.transpose(x, axes=[0, 2, 3, 1])  # To NHWC
          y = relay.layout_transform(y, "NHWC", "HWCN")  # To HWCN
          y = relay.transpose(y, axes=[3, 0, 1, 2])  # To NHWC
          return relay.Function([x], y)
  
      def expected1():
          x = relay.var(
              "x", shape=(1, 3, 224, 224), dtype="float32"
          )  # NCHW
          y = relay.transpose(x, axes=[0, 2, 3, 1])  # To NHWC
          return relay.Function([x], y)
  
      # Test that all transpose ops can be cancelled
      def before2():
          x = relay.var(
              "x", shape=(1, 3, 224, 224), dtype="float32"
          )  # NCHW
          y = relay.nn.relu(x)
          y = relay.transpose(y, axes=[0, 2, 3, 1])  # To NHWC
          y = relay.transpose(y, axes=[1, 2, 3, 0])  # To HWCN
          y = relay.transpose(y, axes=[3, 2, 0, 1])  # To NCHW
          return relay.Function([x], y)
  
      def expected2():
          x = relay.var(
              "x", shape=(1, 3, 224, 224), dtype="float32"
          )  # NCHW
          y = relay.nn.relu(x)
          return relay.Function([x], y)
  
      # Test default axis (reverse) and negative axis
      def before3():
          x = relay.var(
              "x", shape=(1, 3, 224, 224), dtype="float32"
          )  # NCHW
          y = relay.nn.relu(x)
          y = relay.transpose(y)  # Reverse
          y = relay.transpose(y)  # Reverse
          y = relay.transpose(y, axes=[0, 2, -1, 1])
          y = relay.transpose(y)  # Reverse
          y = relay.transpose(y)  # Reverse
          return relay.Function([x], y)
  
      def expected3():
          x = relay.var(
              "x", shape=(1, 3, 224, 224), dtype="float32"
          )  # NCHW
          y = relay.nn.relu(x)
          y = relay.transpose(y, axes=[0, 2, 3, 1])
          return relay.Function([x], y)
  
      # Test a series of transpose and rank changing layout_transform
      def before4():
          """
          Simplify transpose->layout_transform and its inverse.
  
          Input:
          NHWC -> NCHW -> NCHW4c -> op -> NCHW4c -> NCHW -> NHWC
  
          Simplified:
          NHWC -> NCHW4c -> op -> NCHW4c -> NHWC
          """
          x = relay.var("x", shape=(1, 56, 56, 128), dtype="float32")
          y = relay.transpose(x, axes=[0, 3, 1, 2])
          y = relay.layout_transform(y, "NCHW", "NCHW4c")
          y = relay.nn.relu(y)
          y = relay.layout_transform(y, "NCHW4c", "NCHW")
          y = relay.transpose(y, axes=[0, 2, 3, 1])
          return relay.Function([x], y)
  
      def expected4():
          x = relay.var(
              "x", shape=(1, 56, 56, 128), dtype="float32"
          )  # NHWC
          y = relay.layout_transform(x, "NHWC", "NCHW4c")  # To NCHW4c
          y = relay.nn.relu(y)
          y = relay.layout_transform(y, "NCHW4c", "NHWC")  # To NHWC
          return relay.Function([x], y)
  
      def before5():
          """
          Simplify layout_transform->layout_transform and its inverse.
  
          Input:
          NHWC -> NCHW -> NCHW4c -> op -> NCHW4c -> NCHW -> NHWC
  
          Simplified:
          NHWC -> NCHW4c -> op -> NCHW4c -> NHWC
          """
          x = relay.var(
              "x", shape=(1, 56, 56, 128), dtype="float32"
          )  # NHWC
          y = relay.layout_transform(x, "NHWC", "NCHW")  # To NCHW
          y = relay.layout_transform(y, "NCHW", "NCHW4c")  # To NCHW4c
          y = relay.nn.relu(y)
          y = relay.layout_transform(y, "NCHW4c", "NCHW")  # To NCHW
          y = relay.layout_transform(y, "NCHW", "NHWC")  # To NHWC
          return relay.Function([x], y)
  
      def expected5():
          x = relay.var(
              "x", shape=(1, 56, 56, 128), dtype="float32"
          )  # NHWC
          y = relay.layout_transform(x, "NHWC", "NCHW4c")  # To NCHW4c
          y = relay.nn.relu(y)
          y = relay.layout_transform(y, "NCHW4c", "NHWC")  # To NHWC
          return relay.Function([x], y)
  
      def before6():
          """
          Remove trivial layout_transform->layout_transform.
  
          Input:
          NCHW -> NHWC -> NCHW -> op
  
          Simplified:
          NHWC -> op
          """
  
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.layout_transform(x, "NCHW", "NHWC")
          y = relay.layout_transform(y, "NHWC", "NCHW")
          y = relay.nn.relu(y)
          return relay.Function([x], y)
  
      def expected6():
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.nn.relu(x)
          return relay.Function([x], y)
  
      def before7():
          """
          Remove trivial layout_transform->layout_transform.
  
          Input:
          NCHW4c -> NCHW8c -> NCHW4c -> op
  
          Simplified:
          NCHW4c -> op
          """
          x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32")
          y = relay.layout_transform(x, "NCHW4c", "NCHW8c")
          y = relay.layout_transform(y, "NCHW8c", "NCHW4c")
          y = relay.nn.relu(y)
          return relay.Function([x], y)
  
      def expected7():
          x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32")
          y = relay.nn.relu(x)
          return relay.Function([x], y)
  
      def before8():
          """
          Simplify layout_transform->layout_transform with rank contraction and expansion
  
          Input:
          NCHW4c -> NCHW -> NCHW8c -> op
  
          Simplified:
          NCHW4c -> NCHW8c -> op
          """
          x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32")
          y = relay.layout_transform(x, "NCHW4c", "NCHW")
          y = relay.layout_transform(y, "NCHW", "NCHW8c")
          y = relay.nn.relu(y)
          return relay.Function([x], y)
  
      def expected8():
          x = relay.var("x", shape=(1, 32, 56, 56, 4), dtype="float32")
          y = relay.layout_transform(x, "NCHW4c", "NCHW8c")
          y = relay.nn.relu(y)
          return relay.Function([x], y)
  
      def before9():
          """
          Remove trivial layout_transform->layout_transform.
  
          Input:
          NCHW -> NCHW4c -> NCHW -> op
  
          Simplified:
          NCHW -> op
          """
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.layout_transform(x, "NCHW", "NCHW4c")
          y = relay.layout_transform(y, "NCHW4c", "NCHW")
          y = relay.nn.relu(y)
          return relay.Function([x], y)
  
      def expected9():
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.nn.relu(x)
          return relay.Function([x], y)
  
      def before10():
          """
          Simplify layout_transform->layout_transform without rank change to transpose.
  
          Input:
          NCHW -> NHWC -> CHWN -> op
  
          Simplified:
          NCHW -> CHWN -> op
          """
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.layout_transform(x, "NCHW", "NHWC")
          y = relay.layout_transform(y, "NHWC", "CHWN")
          y = relay.nn.relu(y)
          return relay.Function([x], y)
  
      def expected10():
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.transpose(x, axes=[1, 2, 3, 0])
          y = relay.nn.relu(y)
          return relay.Function([x], y)
  
      def before11():
          """
          Remove trivial no op transpose ops
  
          Input:
          op1 -> relay.transpose(x, axes=[0, 1, 2, 3]) -> op2
  
          Simplified:
          op1 -> op2
          """
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.transpose(x, axes=[0, 1, 2, 3])
          y = relay.nn.relu(y)
          y = relay.layout_transform(y, "NCHW", "NCHW")
          return relay.Function([x], y)
  
      def expected11():
          x = relay.var("x", shape=(1, 128, 56, 56), dtype="float32")
          y = relay.nn.relu(x)
          return relay.Function([x], y)
  
      for before, expected in [
          [before1(), expected1()],
          [before2(), expected2()],
          [before3(), expected3()],
          [before4(), expected4()],
          [before5(), expected5()],
          [before6(), expected6()],
          [before7(), expected7()],
          [before8(), expected8()],
          [before9(), expected9()],
          [before10(), expected10()],
          [before11(), expected11()],
      ]:
          after = run_opt_pass(before, transform.SimplifyExpr())
          expected = run_opt_pass(expected, transform.InferType())
          assert tvm.ir.structural_equal(
              after, expected
          ), "\nafter: {} \nexpected: {}".format(after, expected)
  
imports: from math import sqrt
  import pytest
  import tvm
  from tvm import relay
  from tvm.relay import transform
  from tvm.relay.testing import run_opt_pass, run_infer_type
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: test_simplify_transpose() and returns: relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay, relay