code: def batch_matmul(lhs, rhs, transa=False, transb=False):
      """Create an extern op that compute matrix mult of A and rhs with rocBLAS
  
      Parameters
      ----------
      lhs : Tensor
          The left batched matrix operand
      rhs : Tensor
          The right batched matrix operand
      transa : bool
          Whether transpose lhs
      transb : bool
          Whether transpose rhs
  
      Returns
      -------
      C : Tensor
          The result tensor.
      """
      batch_size = lhs.shape[0]
      assert batch_size == rhs.shape[0]
      n = lhs.shape[2] if transa else lhs.shape[1]
      m = rhs.shape[1] if transb else rhs.shape[2]
      return te.extern(
          (batch_size, n, m),
          [lhs, rhs],
          lambda ins, outs: tvm.tir.call_packed(
              "tvm.contrib.rocblas.batch_matmul",
              ins[0],
              ins[1],
              outs[0],
              transa,
              transb,
          ),
          name="C",
      )
  
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) and returns: te