Repository: dmlc/dgl

Issue:
[Sparse] Support SparseMatrix element-wise multiplication with different sparsities.
## 🔨Work Item

**IMPORTANT:**
* This template is only for dev team to track project progress. For feature request or bug report, please use the corresponding issue templates.
* DO NOT create a new work item if the purpose is to fix an existing issue or feature request. We will directly use the issue in the project tracker.

Project tracker: https://github.com/orgs/dmlc/projects/2

## Description

<del>
Breakdown into two work items.

1. Support on CPU (#5519).

3. Support on GPU.

</del>

Implemented by concatenation and unique on two COO matrices.
## Depending work items or issues

<!-- what must be done before this -->



Discussion:

Resolving patch:
diff --git a/python/dgl/sparse/elementwise_op_sp.py b/python/dgl/sparse/elementwise_op_sp.py
--- a/python/dgl/sparse/elementwise_op_sp.py
+++ b/python/dgl/sparse/elementwise_op_sp.py
@@ -14,6 +14,13 @@ def spsp_add(A, B):
     )
 
 
+def spsp_mul(A, B):
+    """Invoke C++ sparse library for multiplication"""
+    return SparseMatrix(
+        torch.ops.dgl_sparse.spsp_mul(A.c_sparse_matrix, B.c_sparse_matrix)
+    )
+
+
 def sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
     """Elementwise addition
 
@@ -83,8 +90,8 @@ def sp_sub(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
 def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
     """Elementwise multiplication
 
-    If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must be
-    diagonal matrices.
+    Note that if both :attr:`A` and :attr:`B` are sparse matrices, both of them
+    need to be diagonal or on CPU.
 
     Parameters
     ----------
@@ -108,28 +115,27 @@ def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
     >>> A * 2
     SparseMatrix(indices=tensor([[1, 0, 2],
                                  [0, 3, 2]]),
-    values=tensor([2, 4, 6]),
-    shape=(3, 4), nnz=3)
+                 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)
+                 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)
-    if A.is_diag() and B.is_diag():
-        assert A.shape == B.shape, (
-            f"The shape of diagonal matrix A {A.shape} and B {B.shape} must"
-            f"match for elementwise multiplication."
-        )
-        return diag(A.val * B.val, A.shape)
-    # Python falls back to B.__rmul__(A) then TypeError when NotImplemented is
-    # returned.
-    # So this also handles the case of scalar * SparseMatrix since we set
-    # SparseMatrix.__rmul__ to be the same as SparseMatrix.__mul__.
-    return NotImplemented
+    return spsp_mul(A, B)
 
 
 def sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:


Test patch:
diff --git a/tests/python/pytorch/sparse/test_elementwise_op.py b/tests/python/pytorch/sparse/test_elementwise_op.py
--- a/tests/python/pytorch/sparse/test_elementwise_op.py
+++ b/tests/python/pytorch/sparse/test_elementwise_op.py
@@ -225,7 +225,7 @@ def test_sub_sparse_diag(val_shape):
     assert torch.allclose(dense_diff, -diff4)
 
 
-@pytest.mark.parametrize("op", ["mul", "truediv", "pow"])
+@pytest.mark.parametrize("op", ["truediv", "pow"])
 def test_error_op_sparse_diag(op):
     ctx = F.ctx()
     row = torch.tensor([1, 0, 2]).to(ctx)
diff --git a/tests/python/pytorch/sparse/test_elementwise_op_sp.py b/tests/python/pytorch/sparse/test_elementwise_op_sp.py
--- a/tests/python/pytorch/sparse/test_elementwise_op_sp.py
+++ b/tests/python/pytorch/sparse/test_elementwise_op_sp.py
@@ -4,7 +4,15 @@
 import pytest
 import torch
 
-from dgl.sparse import from_coo, power
+from dgl.sparse import from_coo, mul, power, val_like
+
+from .utils import (
+    rand_coo,
+    rand_csc,
+    rand_csr,
+    rand_diag,
+    sparse_matrix_to_dense,
+)
 
 
 def all_close_sparse(A, row, col, val, shape):
@@ -91,3 +99,38 @@ def test_error_op_scalar(op, v_scalar):
         A - v_scalar
     with pytest.raises(TypeError):
         v_scalar - A
+
+
+@pytest.mark.parametrize(
+    "create_func1", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize(
+    "create_func2", [rand_coo, rand_csr, rand_csc, rand_diag]
+)
+@pytest.mark.parametrize("shape", [(5, 5), (5, 3)])
+@pytest.mark.parametrize("nnz1", [5, 15])
+@pytest.mark.parametrize("nnz2", [1, 14])
+@pytest.mark.parametrize("nz_dim", [None, 3])
+def test_spspmul(create_func1, create_func2, shape, nnz1, nnz2, nz_dim):
+    dev = F.ctx()
+    A = create_func1(shape, nnz1, dev, nz_dim)
+    B = create_func2(shape, nnz2, dev, nz_dim)
+    C = mul(A, B)
+    assert not C.has_duplicate()
+
+    DA = sparse_matrix_to_dense(A)
+    DB = sparse_matrix_to_dense(B)
+    DC = DA * DB
+
+    grad = torch.rand_like(C.val)
+    C.val.backward(grad)
+    DC_grad = sparse_matrix_to_dense(val_like(C, grad))
+    DC.backward(DC_grad)
+
+    assert torch.allclose(sparse_matrix_to_dense(C), DC, atol=1e-05)
+    assert torch.allclose(
+        val_like(A, A.val.grad).to_dense(), DA.grad, atol=1e-05
+    )
+    assert torch.allclose(
+        val_like(B, B.val.grad).to_dense(), DB.grad, atol=1e-05
+    )
diff --git a/tests/python/pytorch/sparse/utils.py b/tests/python/pytorch/sparse/utils.py
--- a/tests/python/pytorch/sparse/utils.py
+++ b/tests/python/pytorch/sparse/utils.py
@@ -1,7 +1,7 @@
 import numpy as np
 import torch
 
-from dgl.sparse import from_csc, from_csr, SparseMatrix, spmatrix
+from dgl.sparse import diag, from_csc, from_csr, SparseMatrix, spmatrix
 
 np.random.seed(42)
 torch.random.manual_seed(42)
@@ -83,6 +83,15 @@ def rand_csc(shape, nnz, dev, nz_dim=None):
     return from_csc(indptr, indices, val, shape=shape)
 
 
+def rand_diag(shape, nnz, dev, nz_dim=None):
+    nnz = min(shape)
+    if nz_dim is None:
+        val = torch.randn(nnz, device=dev, requires_grad=True)
+    else:
+        val = torch.randn(nnz, nz_dim, device=dev, requires_grad=True)
+    return diag(val, shape)
+
+
 def rand_coo_uncoalesced(shape, nnz, dev):
     # Create a sparse matrix with possible duplicate entries.
     row = torch.randint(shape[0], (nnz,), device=dev)