code: def batch_matmul(lhs, rhs, transa=False, transb=False, dtype=None):
      """Create an extern op that compute batch matrix mult of A and rhs with cuBLAS
  
      Parameters
      ----------
      lhs : Tensor
          The left matrix operand
      rhs : Tensor
          The right matrix operand
      transa : bool
          Whether transpose lhs
      transb : bool
          Whether transpose rhs
  
      Returns
      -------
      C : Tensor
          The result tensor.
      """
      b = lhs.shape[0]
      n = lhs.shape[2] if transa else lhs.shape[1]
      m = rhs.shape[1] if transb else rhs.shape[2]
      dtype = dtype if dtype is not None else lhs.dtype
      return te.extern(
          (b, n, m),
          [lhs, rhs],
          lambda ins, outs: tvm.tir.call_packed(
              "tvm.contrib.cublas.batch_matmul",
              ins[0],
              ins[1],
              outs[0],
              transa,
              transb,
          ),
          dtype=dtype,
          name="batch_matmul_cublas",
      )
  
imports: import tvm
  from tvm import te
label_desc: Use image node_id 2 for calling a global function with example usage: batch_matmul(lhs, rhs, transa, transb, dtype) and returns: te