code: def test_bn_convert_layout():
      def before():
          x = relay.var("x", shape=(1, 56, 56, 64))
          weight1 = relay.var("weight1", shape=(3, 3, 64, 32))
          y = relay.nn.conv2d(
              x,
              weight1,
              channels=32,
              kernel_size=(3, 3),
              padding=(1, 1),
              data_layout="NHWC",
              kernel_layout="HWIO",
          )
          gamma = relay.var("gamma")
          beta = relay.var("beta")
          mean = relay.var("mean")
          variance = relay.var("variance")
          y, _, _ = relay.nn.batch_norm(
              y, gamma, beta, mean, variance, axis=3
          )
          return relay.Function(analysis.free_vars(y), y)
  
      a = before()
      a = run_opt_pass(
          a, transform.ConvertLayout({"nn.conv2d": ["NCHW", "default"]})
      )
  
      # Check that there is only 1 NHWC to NCHW transform.
      has_lt = list()
      find_op = lambda x: has_lt.append(
          isinstance(x, tvm.relay.expr.Call)
          and x.op.name == "layout_transform"
          and x.attrs.src_layout == "NCHW"
          and x.attrs.dst_layout == "NHWC"
      )
      relay.analysis.post_order_visit(a, find_op)
      has_lt = list(filter(lambda x: x, has_lt))
      assert len(has_lt) == 1
  
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 13 for calling a global function with example usage: test_bn_convert_layout() and returns: relay