code: def matmul(
      tensor_a,
      tensor_b,
      bias=None,
      out_dtype=None,
      transpose_a=False,
      transpose_b=False,
      auto_scheduler_rewritten_layout="",
      meta_schedule_original_shape=None,
  ):
      """The default implementation of matmul in topi.
  
      Parameters
      ----------
      tensor_a : tvm.te.Tensor
          2-D with shape [batch, in_dim]
  
      tensor_b : tvm.te.Tensor
          2-D with shape [out_dim, in_dim]
  
      bias : Optional[tvm.te.Tensor]
          1-D with shape [out_dim]
  
      out_dtype : Optional[str]
          The output type. This is used for mixed precision.
  
      transpose_a : Optional[bool] = False
          Whether the tensor_a is in transposed format.
  
      transpose_b : Optional[bool] = False
          Whether the tensor_b is in transposed format.
  
      auto_scheduler_rewritten_layout: Optional[str] = ""
          The layout after auto-scheduler's layout rewrite pass.
  
      meta_schedule_original_shape: Optional[List[PrimExpr]] = None
          The original shape of the input tensor.
  
      Returns
      -------
      output : tvm.te.Tensor
          2-D with shape [batch, out_dim]
      """
      # TODO(yixin): support cases for 1-dim input
      # TODO(yixin): adding support and further check for >2-dim input in autotvm template
      assert (
          len(tensor_a.shape) >= 2 and len(tensor_b.shape) >= 2
      ), "1-dim matmul is not supported yet."
      if bias is not None:
          assert len(bias.shape) == 1
      if out_dtype is None:
          out_dtype = tensor_a.dtype
      if transpose_a:
          reduce_dim_a, in_dim = tensor_a.shape[-2:]
      else:
          in_dim, reduce_dim_a = tensor_a.shape[-2:]
      batch_dims_a = tensor_a.shape[:-2]
  
      if auto_scheduler_rewritten_layout:
          # Infer shape for the rewritten layout
          assert (
              len(tensor_b).shape == 2
          ), "only support 2-dim matmul when using auto-scheduler"
          (
              out_dim,
              reduce_dim_b,
          ) = auto_scheduler.get_shape_from_rewritten_layout(
              auto_scheduler_rewritten_layout, ["j", "k"]
          )
          auto_scheduler.remove_index_check(tensor_b)
      elif meta_schedule_original_shape:
          auto_scheduler.rewrite_tensor_shape(
              tensor_b, meta_schedule_original_shape
          )
          if transpose_b:
              out_dim, reduce_dim_b = tensor_b.shape[-2:]
          else:
              reduce_dim_b, out_dim = tensor_b.shape[-2:]
      elif transpose_b:
          out_dim, reduce_dim_b = tensor_b.shape[-2:]
      else:
          reduce_dim_b, out_dim = tensor_b.shape[-2:]
      batch_dims_b = tensor_b.shape[:-2]
  
      if not isinstance(reduce_dim_a, tvm.tir.Var) and not isinstance(
          reduce_dim_b, tvm.tir.Var
      ):
          assert int(reduce_dim_a) == int(
              reduce_dim_b
          ), f"Reduction dimensions of dense do not match. {reduce_dim_a} vs {reduce_dim_b}."
  
      result_ndim = max(len(batch_dims_a), len(batch_dims_b))
      batch_dims_a = [1] * (
          result_ndim - len(batch_dims_a)
      ) + batch_dims_a
      batch_dims_b = [1] * (
          result_ndim - len(batch_dims_b)
      ) + batch_dims_b
  
      for idx, (l, r) in enumerate(zip(batch_dims_a, batch_dims_b)):
          if (
              not isinstance(l, tvm.tir.Var)
              and not isinstance(r, tvm.tir.Var)
              and int(l) != 1
              and int(r) != 1
          ):
              assert int(l) == int(r), (
                  "Batch dimensions of dense do not match: "
                  f"{tensor_a.shape[:-2]} vs {tensor_b.shape[:-2]}."
              )
          if not isinstance(l, tvm.tir.Var) and int(l) == 1:
              batch_dims_a[idx] = batch_dims_b[idx]
  
      k = te.reduce_axis((0, reduce_dim_a), name="k")
  
      def compute(*indices):
          batch_indices_a = indices[-len(tensor_a.shape) : -2]
          batch_indices_a = [
              i if isinstance(dim, tvm.tir.Var) or int(dim) != 1 else 0
              for i, dim in zip(batch_indices_a, tensor_a.shape[:-2])
          ]
          batch_indices_b = indices[-len(tensor_b.shape) : -2]
          batch_indices_b = [
              i if isinstance(dim, tvm.tir.Var) or int(dim) != 1 else 0
              for i, dim in zip(batch_indices_b, tensor_b.shape[:-2])
          ]
          i, j = indices[-2:]
          a_indices = (
              (*batch_indices_a, k, i)
              if transpose_a
              else (*batch_indices_a, i, k)
          )
          b_indices = (
              (*batch_indices_b, j, k)
              if transpose_b
              else (*batch_indices_b, k, j)
          )
          return te.sum(
              tensor_a[a_indices].astype(out_dtype)
              * tensor_b[b_indices].astype(out_dtype),
              axis=k,
          )
  
      compute_name = {
          (True, True): "T_matmul_TT",
          (True, False): "T_matmul_TN",
          (False, True): "T_matmul_NT",
          (False, False): "T_matmul_NN",
      }[(transpose_a, transpose_b)]
  
      # TODO(jcf94): Remove `dense` when `matmul` is finally ready
      compute_tag = (
          "dense"
          if (transpose_a, transpose_b) == (False, True)
          else "matmul"
      )
  
      mat = te.compute(
          (*batch_dims_a, in_dim, out_dim),
          compute,
          name=compute_name,
          tag=compute_tag,
          attrs={"layout_free_placeholders": [tensor_b]},
      )
  
      if bias is not None:
          mat = add(mat, bias.astype(out_dtype))
  
      if auto_scheduler_rewritten_layout:
          mat = auto_scheduler.rewrite_compute_body(
              mat, auto_scheduler_rewritten_layout
          )
  
      return mat
  
file_path: tvm/python/tvm/topi/nn/dense.py
imports: import tvm
  from tvm import auto_scheduler, te
  from ..None import tag, add
label_desc: Use image node_id 1 for calling a global function with example usage: matmul(tensor_a, tensor_b, bias, out_dtype, transpose_a, transpose_b, auto_scheduler_rewritten_layout, meta_schedule_original_shape) and returns: mat, te