code: def test_infer_correct_layout():
      test_infer_correct_layout_flag = False
  
      def before():
          x = relay.var("x", shape=(1, 56, 56, 64))
          weight = relay.var("weight", shape=(3, 3, 64, 64))
          y = relay.nn.conv2d(
              x,
              weight,
              channels=64,
              kernel_size=(3, 3),
              padding=(1, 1),
              data_layout="NHWC",
              kernel_layout="HWIO",
          )
          y = relay.nn.relu(y)
          y = relay.Function([x, weight], y)
          return y
  
      @reg.register_infer_correct_layout("nn.relu", level=11)
      def infer_correct_layout_relu(
          attrs, new_in_layouts, old_in_layouts, old_in_types
      ):
          nonlocal test_infer_correct_layout_flag
          test_infer_correct_layout_flag = True
          ret = tvm.tir.layout("")
          if new_in_layouts:
              assert len(new_in_layouts) >= 1
              ret = new_in_layouts[0]
          else:
              for i in range(len(old_in_layouts)):
                  if old_in_layouts[i]:
                      ret = old_in_layouts[i]
                      break
          input_layouts = []
          for i in range(len(old_in_layouts)):
              input_layouts.append(ret)
          return InferCorrectLayoutOutput(input_layouts, [ret], attrs)
  
      a = before()
      a = run_opt_pass(
          a, transform.ConvertLayout({"nn.conv2d": ["NCHW", "default"]})
      )
      assert test_infer_correct_layout_flag == True
  
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 42 for calling a global function with example usage: test_infer_correct_layout() and returns: y, InferCorrectLayoutOutput