code: def verify_batch_matmul(
      x_batch, y_batch, M, N, K, dynamic=False, debug=False
  ):
      if not dynamic:
          x = te.placeholder((x_batch, M, K), name="x")
          y = te.placeholder((y_batch, N, K), name="y")
          dtype = x.dtype
      else:
          assert x_batch == y_batch or x_batch == 1 or y_batch == 1
          batch_size = max(x_batch, y_batch)
          dynamic_batch_size = te.var("dynamic_batch_size")
          dynamic_M = te.var("dynamic_M")
          dynamic_N = te.var("dynamic_N")
          dynamic_K = te.var("dynamic_K")
  
          x = te.placeholder(
              (dynamic_batch_size, dynamic_M, dynamic_K), name="x"
          )
          y = te.placeholder(
              (dynamic_batch_size, dynamic_N, dynamic_K), name="y"
          )
          dtype = x.dtype
  
      # use memoize to pickle the test data for next time use
      @memoize("topi.tests.test_topi_batch_matmul")
      def get_ref_data():
          a_np = np.random.uniform(size=(x_batch, M, K)).astype(dtype)
          b_np = np.random.uniform(size=(y_batch, N, K)).astype(dtype)
          c_np = tvm.topi.testing.batch_matmul(a_np, b_np)
          return (a_np, b_np, c_np)
  
      # get the test data
      a_np, b_np, c_np = get_ref_data()
  
      def check_device(target, dev):
          print("Running on target: %s" % target)
          with tvm.target.Target(target):
              fcompute, fschedule = tvm.topi.testing.dispatch(
                  target, _batch_matmul_implement
              )
              out = fcompute(x, y)
              if not dynamic:
                  s = fschedule([out])
                  out_shape = out.shape
              else:
                  s = te.create_schedule(out.op)
                  out_shape = (batch_size, M, N)
  
              if debug:
                  print(tvm.lower(s, [x, y, out], simple_mode=True))
  
          a = tvm.nd.array(a_np, dev)
          b = tvm.nd.array(b_np, dev)
          c = tvm.nd.array(
              np.zeros(get_const_tuple(out_shape), dtype=dtype), dev
          )
          f = tvm.build(s, [x, y, out], target, name="dense")
          f(a, b, c)
          tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5)
  
      for target, dev in tvm.testing.enabled_targets():
          target_kind = tvm.target.Target(target).kind.name
          if dynamic and target_kind in [
              "cuda",
              "nvptx",
              "vulkan",
              "opencl",
          ]:
              print(
                  "Dynamic batch matmul test is skippped on %s" % target
              )
              continue
  
          check_device(target, dev)
  
file_path: tvm/tests/python/topi/test_topi_batch_matmul.py
imports: import numpy
  import tvm
  from tvm import te
  from tvm import topi
  import tvm.topi.testing
  from tvm.topi.utils import get_const_tuple
  from tvm.contrib.pickle_memoize import memoize
  import tvm.testing
  from common import Int8Fallback
label_desc: Use image node_id 1 for calling a global function with example usage: verify_batch_matmul(x_batch, y_batch, M, N, K, dynamic, debug) and returns: a_np,  b_np,  c_np