code: def matmul(
      A: Union[torch.Tensor, SparseMatrix],
      B: Union[torch.Tensor, SparseMatrix],
  ) -> Union[torch.Tensor, SparseMatrix]:
      """Multiplies two dense/sparse matrices, equivalent to ``A @ B``.
  
      This function does not support the case where :attr:`A` is a \
      ``torch.Tensor`` and :attr:`B` is a ``SparseMatrix``.
  
      * If both matrices are torch.Tensor, it calls \
          :func:`torch.matmul()`. The result is a dense matrix.
  
      * If both matrices are sparse, it calls :func:`dgl.sparse.spspmm`. The \
          result is a sparse matrix.
  
      * If :attr:`A` is sparse while :attr:`B` is dense, it calls \
          :func:`dgl.sparse.spmm`. The result is a dense matrix.
  
      * The operator supports batched sparse-dense matrix multiplication. In \
          this case, the sparse matrix :attr:`A` should have shape ``(L, M)``, \
          where the non-zero values have a batch dimension ``K``. The dense \
          matrix :attr:`B` should have shape ``(M, N, K)``. The output \
          is a dense matrix of shape ``(L, N, K)``.
  
      * Sparse-sparse matrix multiplication does not support batched computation.
  
      Parameters
      ----------
      A : torch.Tensor or SparseMatrix
          The first matrix.
      B : torch.Tensor or SparseMatrix
          The second matrix.
  
      Returns
      -------
      torch.Tensor or SparseMatrix
          The result matrix
  
      Examples
      --------
  
      Multiplies a diagonal matrix with a dense matrix.
  
      >>> val = torch.randn(3)
      >>> A = dglsp.diag(val)
      >>> B = torch.randn(3, 2)
      >>> result = dglsp.matmul(A, B)
      >>> type(result)
      <class 'torch.Tensor'>
      >>> result.shape
      torch.Size([3, 2])
  
      Multiplies a sparse matrix with a dense matrix.
  
      >>> indices = torch.tensor([[0, 1, 1], [1, 0, 1]])
      >>> val = torch.randn(indices.shape[1])
      >>> A = dglsp.spmatrix(indices, val)
      >>> X = torch.randn(2, 3)
      >>> result = dglsp.matmul(A, X)
      >>> type(result)
      <class 'torch.Tensor'>
      >>> result.shape
      torch.Size([2, 3])
  
      Multiplies a sparse matrix with a sparse matrix.
  
      >>> indices1 = torch.tensor([[0, 1, 1], [1, 0, 1]])
      >>> val1 = torch.ones(indices1.shape[1])
      >>> A = dglsp.spmatrix(indices1, val1)
      >>> indices2 = torch.tensor([[0, 1, 1], [0, 2, 1]])
      >>> val2 = torch.ones(indices2.shape[1])
      >>> B = dglsp.spmatrix(indices2, val2)
      >>> result = dglsp.matmul(A, B)
      >>> type(result)
      <class 'dgl.sparse.sparse_matrix.SparseMatrix'>
      >>> result.shape
      (2, 3)
      """
      assert isinstance(
          A, (torch.Tensor, SparseMatrix)
      ), f"Expect arg1 to be a torch.Tensor or SparseMatrix, got {type(A)}."
      assert isinstance(B, (torch.Tensor, SparseMatrix)), (
          f"Expect arg2 to be a torch Tensor or SparseMatrix"
          f"object, got {type(B)}."
      )
      if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
          return torch.matmul(A, B)
      assert not isinstance(A, torch.Tensor), (
          f"Expect arg2 to be a torch Tensor if arg 1 is torch Tensor, "
          f"got {type(B)}."
      )
      if isinstance(B, torch.Tensor):
          return spmm(A, B)
      return spspmm(A, B)
  
imports: from typing import Union
  import torch
  from .sparse_matrix import SparseMatrix
label_desc: Use image node_id 4 for calling a global function with example usage: matmul(A, B) and returns: spspmm, torch, spmm