def mul(
    A: Union[SparseMatrix, Scalar], B: Union[SparseMatrix, Scalar]
) -> SparseMatrix:
    r"""Elementwise multiplication for ``SparseMatrix``, equivalent to
    ``A * B``.

    If both :attr:`A` and :attr:`B` are sparse matrices, both of them should be
    diagonal matrices.

    Parameters
    ----------
    A : SparseMatrix or Scalar
        Sparse matrix or scalar value
    B : SparseMatrix or Scalar
        Sparse matrix or scalar value

    Returns
    -------
    SparseMatrix
        Sparse matrix

    Examples
    --------
    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
    >>> val = torch.tensor([10, 20, 30])
    >>> A = dglsp.spmatrix(indices, val)
    >>> dglsp.mul(A, 2)
    SparseMatrix(indices=tensor([[1, 0, 2],
                                 [0, 3, 2]]),
                 values=tensor([20, 40, 60]),
                 shape=(3, 4), nnz=3)

    >>> D = dglsp.diag(torch.arange(1, 4))
    >>> dglsp.mul(D, 2)
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([2, 4, 6]),
                 shape=(3, 3), nnz=3)

    >>> D = dglsp.diag(torch.arange(1, 4))
    >>> dglsp.mul(D, D)
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1, 4, 9]),
                 shape=(3, 3), nnz=3)
    """
    return A * B
Elementwise multiplication for ``SparseMatrix``, equivalent to
    ``A * B``.

    If both :attr:`A` and :attr:`B` are sparse matrices, both of them should be
    diagonal matrices.

    Parameters
    ----------
    A : SparseMatrix or Scalar
        Sparse matrix or scalar value
    B : SparseMatrix or Scalar
        Sparse matrix or scalar value

    Returns
    -------
    SparseMatrix
        Sparse matrix

    Examples
    --------
    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
    >>> val = torch.tensor([10, 20, 30])
    >>> A = dglsp.spmatrix(indices, val)
    >>> dglsp.mul(A, 2)
    SparseMatrix(indices=tensor([[1, 0, 2],
                                 [0, 3, 2]]),
                 values=tensor([20, 40, 60]),
                 shape=(3, 4), nnz=3)

    >>> D = dglsp.diag(torch.arange(1, 4))
    >>> dglsp.mul(D, 2)
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([2, 4, 6]),
                 shape=(3, 3), nnz=3)

    >>> D = dglsp.diag(torch.arange(1, 4))
    >>> dglsp.mul(D, D)
    SparseMatrix(indices=tensor([[0, 1, 2],
                                 [0, 1, 2]]),
                 values=tensor([1, 4, 9]),
                 shape=(3, 3), nnz=3)
    
python/dgl/sparse/elementwise_op.py