code: def test_conv_squeeze_convert_layout():
      def _test_conv_squeeze_convert_layout1():
          # specified axis is squeezed
          def before():
              x = relay.var("x", shape=(1, 1, 1, 2048))
              weight = relay.var("weight", shape=(1, 1, 2048, 1000))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=1000,
                  kernel_size=(1, 1),
                  data_layout="NHWC",
                  kernel_layout="HWIO",
              )
              y = relay.nn.relu(y)
              y = relay.squeeze(y, axis=[-3])
              return relay.Function(analysis.free_vars(y), y)
  
          def expected():
              x = relay.var("x", shape=(1, 1, 1, 2048))
              weight = relay.var("weight", shape=(1, 1, 2048, 1000))
              weight = relay.layout_transform(weight, "HWIO", "OIHW")
              x = relay.layout_transform(x, "NHWC", "NCHW")
              y = relay.nn.conv2d(
                  x, weight, channels=1000, kernel_size=(1, 1)
              )
              y = relay.nn.relu(y)
              y = relay.squeeze(y, axis=[2])
              y = relay.layout_transform(y, "NCW", "NWC")
              return relay.Function(analysis.free_vars(y), y)
  
          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_squeeze_convert_layout2():
          # all axes of dimension 1 are squeezed
          def before():
              x = relay.var("x", shape=(1, 1, 1, 2048))
              weight = relay.var("weight", shape=(1, 1, 2048, 1000))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=1000,
                  kernel_size=(1, 1),
                  data_layout="NHWC",
                  kernel_layout="HWIO",
              )
              y = relay.nn.relu(y)
              y = relay.squeeze(y)
              return relay.Function(analysis.free_vars(y), y)
  
          def expected():
              x = relay.var("x", shape=(1, 1, 1, 2048))
              weight = relay.var("weight", shape=(1, 1, 2048, 1000))
              weight = relay.layout_transform(weight, "HWIO", "OIHW")
              x = relay.layout_transform(x, "NHWC", "NCHW")
              y = relay.nn.conv2d(
                  x, weight, channels=1000, kernel_size=(1, 1)
              )
              y = relay.nn.relu(y)
              y = relay.squeeze(y, [0, 2, 3])
              return relay.Function(analysis.free_vars(y), y)
  
          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_squeeze_convert_layout3():
          # squeeze axis is empty
          def before():
              x = relay.var("x", shape=(1, 1, 1, 2048))
              weight = relay.var("weight", shape=(1, 1, 2048, 1000))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=1000,
                  kernel_size=(1, 1),
                  data_layout="NHWC",
                  kernel_layout="HWIO",
              )
              y = relay.nn.relu(y)
              y = relay.squeeze(y, axis=[])
              return relay.Function(analysis.free_vars(y), y)
  
          def expected():
              x = relay.var("x", shape=(1, 1, 1, 2048))
              weight = relay.var("weight", shape=(1, 1, 2048, 1000))
              weight = relay.layout_transform(weight, "HWIO", "OIHW")
              x = relay.layout_transform(x, "NHWC", "NCHW")
              y = relay.nn.conv2d(
                  x, weight, channels=1000, kernel_size=(1, 1)
              )
              y = relay.nn.relu(y)
              y = relay.squeeze(y, axis=[])
              y = relay.layout_transform(y, "NCHW", "NHWC")
              return relay.Function(analysis.free_vars(y), y)
  
          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)
  
      _test_conv_squeeze_convert_layout1()
      _test_conv_squeeze_convert_layout2()
      _test_conv_squeeze_convert_layout3()
  
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 38 for calling a global function with example usage: test_conv_squeeze_convert_layout() and returns: relay, relay, relay, relay, relay, relay