Repository: dmlc/dgl

Issue:
[Sparse] coalesce sparse matrix
## 🔨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

<!-- short description of the work item -->

* Add methods to `SparseMatrix` for coalescing and checking duplicate indices.
* Add duplicate check to spspmm.



Discussion:

Resolving patch:
diff --git a/python/dgl/mock_sparse2/sparse_matrix.py b/python/dgl/mock_sparse2/sparse_matrix.py
--- a/python/dgl/mock_sparse2/sparse_matrix.py
+++ b/python/dgl/mock_sparse2/sparse_matrix.py
@@ -66,6 +66,28 @@ def device(self) -> torch.device:
         """
         return self.c_sparse_matrix.device()
 
+    @property
+    def row(self) -> torch.Tensor:
+        """Get the row indices of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Row indices of the nonzero elements
+        """
+        return self.coo()[0]
+
+    @property
+    def col(self) -> torch.Tensor:
+        """Get the column indices of the nonzero elements.
+
+        Returns
+        -------
+        tensor
+            Column indices of the nonzero elements
+        """
+        return self.coo()[1]
+
     def indices(
         self, fmt: str, return_shuffle=False
     ) -> Tuple[torch.Tensor, ...]:
@@ -173,6 +195,60 @@ def transpose(self):
         """
         return SparseMatrix(self.c_sparse_matrix.transpose())
 
+    def coalesce(self):
+        """Return a coalesced sparse matrix.
+
+        A coalesced sparse matrix satisfies the following properties:
+
+          - the indices of the non-zero elements are unique,
+          - the indices are sorted in lexicographical order.
+
+        The coalescing process will accumulate the non-zero values of the same
+        indices by summation.
+
+        The function does not support autograd.
+
+        Returns
+        -------
+        SparseMatrix
+            The coalesced sparse matrix.
+
+        Examples
+        --------
+        >>> row = torch.tensor([1, 0, 0, 0, 1])
+        >>> col = torch.tensor([1, 1, 1, 2, 2])
+        >>> val = torch.tensor([0, 1, 2, 3, 4])
+        >>> A = create_from_coo(row, col, val)
+        >>> A = A.coalesce()
+        >>> print(A)
+        SparseMatrix(indices=tensor([[0, 0, 1, 1],
+                [1, 2, 1, 2]]),
+        values=tensor([3, 3, 0, 4]),
+        shape=(2, 3), nnz=4)
+        """
+        return SparseMatrix(self.c_sparse_matrix.coalesce())
+
+    def has_duplicate(self):
+        """Return whether this sparse matrix contains duplicate indices.
+
+        Returns
+        -------
+        bool
+            True if this sparse matrix contains duplicate indices.
+
+        Examples
+        --------
+        >>> row = torch.tensor([1, 0, 0, 0, 1])
+        >>> col = torch.tensor([1, 1, 1, 2, 2])
+        >>> val = torch.tensor([0, 1, 2, 3, 4])
+        >>> A = create_from_coo(row, col, val)
+        >>> print(A.has_duplicate())
+        True
+        >>> print(A.coalesce().has_duplicate())
+        False
+        """
+        return self.c_sparse_matrix.has_duplicate()
+
 
 def create_from_coo(
     row: torch.Tensor,


Test patch:
diff --git a/tests/pytorch/mock_sparse2/test_matmul.py b/tests/pytorch/mock_sparse2/test_matmul.py
--- a/tests/pytorch/mock_sparse2/test_matmul.py
+++ b/tests/pytorch/mock_sparse2/test_matmul.py
@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from dgl.mock_sparse2 import val_like
+from dgl.mock_sparse2 import create_from_coo, val_like
 
 from .utils import (
     clone_detach_and_grad,
@@ -59,9 +59,7 @@ def test_spmm(create_func, shape, nnz, out_dim):
 @pytest.mark.parametrize("shape_k", [3, 4])
 @pytest.mark.parametrize("nnz1", [1, 10])
 @pytest.mark.parametrize("nnz2", [1, 10])
-def test_sparse_sparse_mm(
-    create_func1, create_func2, shape_n_m, shape_k, nnz1, nnz2
-):
+def test_spspmm(create_func1, create_func2, shape_n_m, shape_k, nnz1, nnz2):
     dev = F.ctx()
     shape1 = shape_n_m
     shape2 = (shape_n_m[1], shape_k)
@@ -89,3 +87,33 @@ def test_sparse_sparse_mm(
             torch_A2.grad.to_dense(),
             atol=1e-05,
         )
+
+
+def test_spspmm_duplicate():
+    dev = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(dev)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(dev)
+    val = torch.randn(len(row)).to(dev)
+    shape = (4, 4)
+    A1 = create_from_coo(row, col, val, shape)
+
+    row = torch.tensor([1, 0, 0, 1]).to(dev)
+    col = torch.tensor([1, 1, 2, 2]).to(dev)
+    val = torch.randn(len(row)).to(dev)
+    shape = (4, 4)
+    A2 = create_from_coo(row, col, val, shape)
+
+    try:
+        A1 @ A2
+    except:
+        pass
+    else:
+        assert False, "Should raise error."
+
+    try:
+        A2 @ A1
+    except:
+        pass
+    else:
+        assert False, "Should raise error."
diff --git a/tests/pytorch/mock_sparse2/test_sparse_matrix.py b/tests/pytorch/mock_sparse2/test_sparse_matrix.py
--- a/tests/pytorch/mock_sparse2/test_sparse_matrix.py
+++ b/tests/pytorch/mock_sparse2/test_sparse_matrix.py
@@ -1,10 +1,15 @@
-import pytest
-import torch
 import sys
 
 import backend as F
+import pytest
+import torch
 
-from dgl.mock_sparse2 import create_from_coo, create_from_csr, create_from_csc, val_like
+from dgl.mock_sparse2 import (
+    create_from_coo,
+    create_from_csc,
+    create_from_csr,
+    val_like,
+)
 
 # TODO(#4818): Skipping tests on win.
 if not sys.platform.startswith("linux"):
@@ -337,6 +342,7 @@ def test_csr_to_csc(dense_dim, indptr, indices, shape):
     assert torch.allclose(mat_indptr, indptr)
     assert torch.allclose(mat_indices, indices)
 
+
 @pytest.mark.parametrize("val_shape", [(3), (3, 2)])
 @pytest.mark.parametrize("shape", [(3, 5), (5, 5)])
 def test_val_like(val_shape, shape):
@@ -368,3 +374,47 @@ def check_val_like(A, B):
     csc_A = create_from_csc(indptr, indices, val, shape)
     csc_B = val_like(csc_A, new_val)
     check_val_like(csc_A, csc_B)
+
+
+def test_coalesce():
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(ctx)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(ctx)
+    val = torch.arange(len(row)).to(ctx)
+    A = create_from_coo(row, col, val, (4, 4))
+
+    assert A.has_duplicate()
+
+    A_coalesced = A.coalesce()
+
+    assert A_coalesced.nnz == 4
+    assert A_coalesced.shape == (4, 4)
+    assert list(A_coalesced.row) == [0, 0, 1, 1]
+    assert list(A_coalesced.col) == [1, 2, 1, 2]
+    # Values of duplicate indices are added together.
+    assert list(A_coalesced.val) == [3, 3, 0, 4]
+    assert not A_coalesced.has_duplicate()
+
+
+def test_has_duplicate():
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 0, 0, 0, 1]).to(ctx)
+    col = torch.tensor([1, 1, 1, 2, 2]).to(ctx)
+    val = torch.arange(len(row)).to(ctx)
+    shape = (4, 4)
+
+    # COO
+    coo_A = create_from_coo(row, col, val, shape)
+    assert coo_A.has_duplicate()
+
+    # CSR
+    indptr, indices, _ = coo_A.csr()
+    csr_A = create_from_csr(indptr, indices, val, shape)
+    assert csr_A.has_duplicate()
+
+    # CSC
+    indptr, indices, _ = coo_A.csc()
+    csc_A = create_from_csc(indptr, indices, val, shape)
+    assert csc_A.has_duplicate()