def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
    """Elementwise multiplication

    Note that if both :attr:`A` and :attr:`B` are sparse matrices, both of them
    need to be diagonal or on CPU.

    Parameters
    ----------
    A : SparseMatrix
        First operand
    B : SparseMatrix or Scalar
        Second operand

    Returns
    -------
    SparseMatrix
        Result of A * B

    Examples
    --------

    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
    >>> val = torch.tensor([1, 2, 3])
    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))

    >>> A * 2
    SparseMatrix(indices=tensor([[1, 0, 2],
                                 [0, 3, 2]]),
                 values=tensor([2, 4, 6]),
                 shape=(3, 4), nnz=3)

    >>> 2 * A
    SparseMatrix(indices=tensor([[1, 0, 2],
                                 [0, 3, 2]]),
                 values=tensor([2, 4, 6]),
                 shape=(3, 4), nnz=3)

    >>> indices2 = torch.tensor([[2, 0, 1], [0, 3, 2]])
    >>> val2 = torch.tensor([3, 2, 1])
    >>> B = dglsp.spmatrix(indices2, val2, shape=(3, 4))
    >>> A * B
    SparseMatrix(indices=tensor([[0],
                                 [3]]),
                 values=tensor([4]),
                 shape=(3, 4), nnz=1)
    """
    if is_scalar(B):
        return val_like(A, A.val * B)
    return spsp_mul(A, B)
Elementwise multiplication

    Note that if both :attr:`A` and :attr:`B` are sparse matrices, both of them
    need to be diagonal or on CPU.

    Parameters
    ----------
    A : SparseMatrix
        First operand
    B : SparseMatrix or Scalar
        Second operand

    Returns
    -------
    SparseMatrix
        Result of A * B

    Examples
    --------

    >>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
    >>> val = torch.tensor([1, 2, 3])
    >>> A = dglsp.spmatrix(indices, val, shape=(3, 4))

    >>> A * 2
    SparseMatrix(indices=tensor([[1, 0, 2],
                                 [0, 3, 2]]),
                 values=tensor([2, 4, 6]),
                 shape=(3, 4), nnz=3)

    >>> 2 * A
    SparseMatrix(indices=tensor([[1, 0, 2],
                                 [0, 3, 2]]),
                 values=tensor([2, 4, 6]),
                 shape=(3, 4), nnz=3)

    >>> indices2 = torch.tensor([[2, 0, 1], [0, 3, 2]])
    >>> val2 = torch.tensor([3, 2, 1])
    >>> B = dglsp.spmatrix(indices2, val2, shape=(3, 4))
    >>> A * B
    SparseMatrix(indices=tensor([[0],
                                 [3]]),
                 values=tensor([4]),
                 shape=(3, 4), nnz=1)
    
python/dgl/sparse/elementwise_op_sp.py