code: def test_reduce_op_convert_layout():
      for reduce_op in [relay.argmax, relay.mean, relay.max]:
  
          def before():
              x = relay.var("x", shape=(1, 64, 56, 56))
              weight = relay.var("weight", shape=(64, 64, 3, 3))
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=64,
                  kernel_size=(3, 3),
                  padding=(1, 1),
                  data_layout="NCHW",
                  kernel_layout="OIHW",
              )
              y = reduce_op(y, axis=[2, 3])
              y = relay.Function([x, weight], y)
              return y
  
          def expected():
              x = relay.var("x", shape=(1, 64, 56, 56))
              weight = relay.var("weight", shape=(64, 64, 3, 3))
              x = relay.layout_transform(x, "NCHW", "NHWC")
              weight = relay.layout_transform(weight, "OIHW", "HWIO")
              y = relay.nn.conv2d(
                  x,
                  weight,
                  channels=64,
                  kernel_size=(3, 3),
                  padding=(1, 1),
                  data_layout="NHWC",
                  kernel_layout="HWIO",
              )
              y = reduce_op(y, axis=[1, 2])
              y = relay.Function(relay.analysis.free_vars(y), y)
              return y
  
          a = before()
          a = run_opt_pass(
              a,
              transform.ConvertLayout(
                  {"nn.conv2d": ["NHWC", "default"]}
              ),
          )
          b = run_opt_pass(expected(), transform.InferType())
  
          assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
  
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 43 for calling a global function with example usage: test_reduce_op_convert_layout() and returns: y, y