code: def test_fold_batch_norm():
      def expected():
          data = relay.var(
              "data", relay.TensorType((1, 3, 224, 224), "float32")
          )
          weight = relay.const(np.zeros((16, 3, 3, 3)))
          bias = relay.const(np.zeros((16, 1, 1)))
          conv = relay.nn.conv2d(
              data=data,
              weight=weight,
              kernel_size=(3, 3),
              channels=16,
              padding=(1, 1),
          )
          add = relay.add(conv, bias)
          return relay.Function(relay.analysis.free_vars(add), add)
  
      remove_bn_pass = tvm.transform.Sequential(
          [
              relay.transform.InferType(),
              relay.transform.SimplifyInference(),
              relay.transform.FoldConstant(),
              relay.transform.FoldScaleAxis(),
          ]
      )
  
      data = relay.var(
          "data", relay.TensorType((1, 3, 224, 224), "float32")
      )
      weight = relay.var("weight")
      bn_gamma = relay.var("bn_gamma")
      bn_beta = relay.var("bn_beta")
      bn_mmean = relay.var("bn_mean")
      bn_mvar = relay.var("bn_var")
  
      conv = relay.nn.conv2d(
          data=data,
          weight=weight,
          kernel_size=(3, 3),
          channels=16,
          padding=(1, 1),
      )
      bn_output = relay.nn.batch_norm(
          conv, bn_gamma, bn_beta, bn_mmean, bn_mvar
      )
  
      def initializer(_, param):
          param = np.zeros(param.shape)
  
      mod, params = create_workload(bn_output[0], initializer)
      mod["main"] = bind_params_by_name(mod["main"], params)
  
      with tvm.transform.PassContext(opt_level=3):
          mod = remove_bn_pass(mod)
  
      expect = run_infer_type(expected())
      tvm.ir.assert_structural_equal(mod["main"], expect)
  
imports: import numpy
  import tvm
  import tvm.testing
  from tvm import relay
  from tvm.relay.backend import Executor
  from tvm.relay import transform
  from tvm.relay.build_module import bind_params_by_name
  from tvm.relay.testing import run_infer_type, create_workload
label_desc: Use image node_id 13 for calling a global function with example usage: test_fold_batch_norm() and returns: relay