code: def test_conv_split_convert_layout():
      def _test_conv_split_convert_layout1():
          def before():
              x = relay.var("x", shape=(1, 38, 38, 512))
              weight = relay.var("weight", shape=(3, 3, 512, 512))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=512,
                  kernel_size=(3, 3),
                  data_layout="NHWC",
                  kernel_layout="HWIO",
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=2, axis=-1
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              out = relay.Tuple([a, b])
              return relay.Function(analysis.free_vars(out), out)
  
          def expected():
              x = relay.var("x", shape=(1, 38, 38, 512))
              weight = relay.var("weight", shape=(3, 3, 512, 512))
              weight = relay.layout_transform(weight, "HWIO", "OIHW")
              x = relay.layout_transform(x, "NHWC", "NCHW")
              y = relay.nn.conv2d(
                  x, weight, channels=512, kernel_size=(3, 3)
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=2, axis=1
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              a = relay.layout_transform(a, "NCHW", "NHWC")
              b = relay.layout_transform(b, "NCHW", "NHWC")
              out = relay.Tuple([a, b])
              return relay.Function(analysis.free_vars(out), out)
  
          a = before()
          a = run_opt_pass(
              a,
              transform.ConvertLayout(
                  {"nn.conv2d": ["NCHW", "default"]}
              ),
          )
          b = run_opt_pass(expected(), transform.InferType())
  
          assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
  
      def _test_conv_split_convert_layout2():
          def before():
              x = relay.var("x", shape=(1, 38, 38, 512))
              weight = relay.var("weight", shape=(3, 3, 512, 512))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=512,
                  kernel_size=(3, 3),
                  data_layout="NHWC",
                  kernel_layout="HWIO",
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=2, axis=3
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              out = relay.Tuple([a, b])
              return relay.Function(analysis.free_vars(out), out)
  
          def expected():
              x = relay.var("x", shape=(1, 38, 38, 512))
              weight = relay.var("weight", shape=(3, 3, 512, 512))
              weight = relay.layout_transform(weight, "HWIO", "OIHW")
              x = relay.layout_transform(x, "NHWC", "NCHW")
              y = relay.nn.conv2d(
                  x, weight, channels=512, kernel_size=(3, 3)
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=2, axis=1
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              a = relay.layout_transform(a, "NCHW", "NHWC")
              b = relay.layout_transform(b, "NCHW", "NHWC")
              out = relay.Tuple([a, b])
              return relay.Function(analysis.free_vars(out), out)
  
          a = before()
          a = run_opt_pass(
              a,
              transform.ConvertLayout(
                  {"nn.conv2d": ["NCHW", "default"]}
              ),
          )
          b = run_opt_pass(expected(), transform.InferType())
  
          assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
  
      def _test_conv_split_convert_layout3():
          def before():
              x = relay.var("x", shape=(1, 38, 38, 512))
              weight = relay.var("weight", shape=(3, 3, 512, 512))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=512,
                  kernel_size=(3, 3),
                  data_layout="NHWC",
                  kernel_layout="HWIO",
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=(5, 10), axis=-1
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              c = relay.TupleGetItem(y, 2)
              out = relay.Tuple([a, b, c])
              return relay.Function(analysis.free_vars(out), out)
  
          def expected():
              x = relay.var("x", shape=(1, 38, 38, 512))
              weight = relay.var("weight", shape=(3, 3, 512, 512))
              weight = relay.layout_transform(weight, "HWIO", "OIHW")
              x = relay.layout_transform(x, "NHWC", "NCHW")
              y = relay.nn.conv2d(
                  x, weight, channels=512, kernel_size=(3, 3)
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=(5, 10), axis=1
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              c = relay.TupleGetItem(y, 2)
              a = relay.layout_transform(a, "NCHW", "NHWC")
              b = relay.layout_transform(b, "NCHW", "NHWC")
              c = relay.layout_transform(c, "NCHW", "NHWC")
              out = relay.Tuple([a, b, c])
              return relay.Function(analysis.free_vars(out), out)
  
          a = before()
          a = run_opt_pass(
              a,
              transform.ConvertLayout(
                  {"nn.conv2d": ["NCHW", "default"]}
              ),
          )
          b = run_opt_pass(expected(), transform.InferType())
  
          assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
  
      def _test_conv_split_convert_layout_blocking():
          def before():
              x = relay.var("x", shape=(1, 512, 38, 38))
              weight = relay.var("weight", shape=(512, 512, 3, 3))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=512,
                  kernel_size=(3, 3),
                  data_layout="NCHW",
                  kernel_layout="OIHW",
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=[256], axis=1
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              out = relay.Tuple([a, b])
              return relay.Function(analysis.free_vars(out), out)
  
          def expected():
              x = relay.var("x", shape=(1, 512, 38, 38))
              weight = relay.var("weight", shape=(512, 512, 3, 3))
              weight = relay.layout_transform(weight, "OIHW", "OIHW4o")
              x = relay.layout_transform(x, "NCHW", "NCHW4c")
              y = relay.op.nn.contrib_conv2d_nchwc(
                  x,
                  weight,
                  channels=512,
                  kernel_size=(3, 3),
                  padding=(0, 0),
                  data_layout="NCHW4c",
                  kernel_layout="OIHW4o",
              )
              y = relay.nn.relu(y)
              y = relay.op.split(
                  y, indices_or_sections=[64], axis=1
              ).astuple()
              a = relay.TupleGetItem(y, 0)
              b = relay.TupleGetItem(y, 1)
              a = relay.layout_transform(a, "NCHW4c", "NCHW")
              b = relay.layout_transform(b, "NCHW4c", "NCHW")
              out = relay.Tuple([a, b])
              return relay.Function(analysis.free_vars(out), out)
  
          a = before()
          a = run_opt_pass(
              a,
              transform.ConvertLayout(
                  {"nn.conv2d": ["NCHW4c", "OIHW4o"]}
              ),
          )
          b = run_opt_pass(expected(), transform.InferType())
  
          assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
  
      _test_conv_split_convert_layout1()
      _test_conv_split_convert_layout2()
      _test_conv_split_convert_layout3()
      _test_conv_split_convert_layout_blocking()
  
imports: import pytest
  import tvm
  from tvm import relay, te
  from tvm.relay import analysis, transform
  from tvm.relay.op import op
  from tvm.relay.op import register_alter_op_layout
  from tvm.relay.quantize._annotate import attach_simulated_quantize, QAnnotateKind
  from tvm.relay.transform.infer_layout_utils import InferCorrectLayoutOutput
label_desc: Use image node_id 30 for calling a global function with example usage: test_conv_split_convert_layout() and returns: relay, relay, relay, relay, relay, relay, relay, relay