Repository: dmlc/dgl

Issue:
[Sparse] Support sparse format conversion.
## 🔨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 -->

## Depending work items or issues

<!-- what must be done before this -->



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
@@ -109,8 +109,8 @@ def coo(self) -> Tuple[torch.Tensor, ...]:
 
         Returns
         -------
-        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-            A tuple of tensors containing row, column coordinates and values.
+        Tuple[torch.Tensor, torch.Tensor]
+            A tuple of tensors containing row and column coordinates.
         """
         return self.c_sparse_matrix.coo()
 
@@ -120,7 +120,7 @@ def csr(self) -> Tuple[torch.Tensor, ...]:
         Returns
         -------
         Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-            A tuple of tensors containing row, column coordinates and values.
+            A tuple of tensors containing row, column coordinates and value indices.
         """
         return self.c_sparse_matrix.csr()
 
@@ -130,7 +130,7 @@ def csc(self) -> Tuple[torch.Tensor, ...]:
         Returns
         -------
         Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
-            A tuple of tensors containing row, column coordinates and values.
+            A tuple of tensors containing row, column coordinates and value indices.
         """
         return self.c_sparse_matrix.csc()
 
@@ -142,7 +142,8 @@ def dense(self) -> torch.Tensor:
         torch.Tensor
             Dense representation of the sparse matrix.
         """
-        row, col, val = self.coo()
+        row, col = self.coo()
+        val = self.val
         shape = self.shape + val.shape[1:]
         mat = torch.zeros(shape, device=self.device)
         mat[row, col] = val
@@ -186,16 +187,16 @@ def create_from_coo(
     >>> A = create_from_coo(dst, src)
     >>> print(A)
     SparseMatrix(indices=tensor([[1, 1, 2],
-                                 [2, 4, 3]]),
-                 values=tensor([1., 1., 1.]),
-                 shape=(3, 5), nnz=3)
+            [2, 4, 3]]),
+    values=tensor([1., 1., 1.]),
+    shape=(3, 5), nnz=3)
     >>> # Specify shape
     >>> A = create_from_coo(dst, src, shape=(5, 5))
     >>> print(A)
     SparseMatrix(indices=tensor([[1, 1, 2],
-                                 [2, 4, 3]]),
-                 values=tensor([1., 1., 1.]),
-                 shape=(5, 5), nnz=3)
+            [2, 4, 3]]),
+    values=tensor([1., 1., 1.]),
+    shape=(5, 5), nnz=3)
 
     Case2: Sparse matrix with scalar/vector values. Following example is with
     vector data.
@@ -203,11 +204,11 @@ def create_from_coo(
     >>> val = torch.tensor([[1., 1.], [2., 2.], [3., 3.]])
     >>> A = create_from_coo(dst, src, val)
     SparseMatrix(indices=tensor([[1, 1, 2],
-                                 [2, 4, 3]]),
-                 values=tensor([[1, 1],
-                                [2, 2],
-                                [3, 3]]),
-                 shape=(3, 5), nnz=3)
+            [2, 4, 3]]),
+    values=tensor([[1., 1.],
+            [2., 2.],
+            [3., 3.]]),
+    shape=(3, 5), nnz=3)
     """
     if shape is None:
         shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
@@ -219,8 +220,6 @@ def create_from_coo(
     )
 
 
-# FIXME: The docstring cannot print A because we cannot print
-# the indices of CSR/CSC
 def create_from_csr(
     indptr: torch.Tensor,
     indices: torch.Tensor,
@@ -269,16 +268,16 @@ def create_from_csr(
     >>> A = create_from_csr(indptr, indices)
     >>> print(A)
     SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-                                 [1, 2, 0, 1, 2]]),
-                 values=tensor([1., 1., 1., 1., 1.]),
-                 shape=(3, 3), nnz=5)
+            [1, 2, 0, 1, 2]]),
+    values=tensor([1., 1., 1., 1., 1.]),
+    shape=(3, 3), nnz=5)
     >>> # Specify shape
-    >>> A = create_from_csr(indptr, indices, shape=(5, 3))
+    >>> A = create_from_csr(indptr, indices, shape=(3, 5))
     >>> print(A)
     SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
             [1, 2, 0, 1, 2]]),
     values=tensor([1., 1., 1., 1., 1.]),
-    shape=(5, 3), nnz=5)
+    shape=(3, 5), nnz=5)
 
     Case2: Sparse matrix with scalar/vector values. Following example is with
     vector data.
@@ -305,8 +304,6 @@ def create_from_csr(
     )
 
 
-# FIXME: The docstring cannot print A because we cannot print
-# the indices of CSR/CSC
 def create_from_csc(
     indptr: torch.Tensor,
     indices: torch.Tensor,
@@ -354,15 +351,15 @@ def create_from_csc(
     >>> indices = torch.tensor([2, 0, 2, 1, 2])
     >>> A = create_from_csc(indptr, indices)
     >>> print(A)
-    SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-                                 [1, 2, 0, 1, 2]]),
-                 values=tensor([1., 1., 1., 1., 1.]),
-                 shape=(3, 3), nnz=5)
+    SparseMatrix(indices=tensor([[2, 0, 2, 1, 2],
+            [0, 1, 1, 2, 2]]),
+    values=tensor([1., 1., 1., 1., 1.]),
+    shape=(3, 3), nnz=5)
     >>> # Specify shape
     >>> A = create_from_csc(indptr, indices, shape=(5, 3))
     >>> print(A)
-    SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-            [1, 2, 0, 1, 2]]),
+    SparseMatrix(indices=tensor([[2, 0, 2, 1, 2],
+            [0, 1, 1, 2, 2]]),
     values=tensor([1., 1., 1., 1., 1.]),
     shape=(5, 3), nnz=5)
 
@@ -372,12 +369,12 @@ def create_from_csc(
     >>> val = torch.tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
     >>> A = create_from_csc(indptr, indices, val)
     >>> print(A)
-    SparseMatrix(indices=tensor([[0, 1, 2, 2, 2],
-            [1, 2, 0, 1, 2]]),
-    values=tensor([[2, 2],
-            [4, 4],
-            [1, 1],
+    SparseMatrix(indices=tensor([[2, 0, 2, 1, 2],
+            [0, 1, 1, 2, 2]]),
+    values=tensor([[1, 1],
+            [2, 2],
             [3, 3],
+            [4, 4],
             [5, 5]]),
     shape=(3, 3), nnz=5)
     """


Test patch:
diff --git a/tests/pytorch/mock_sparse2/test_diag.py b/tests/pytorch/mock_sparse2/test_diag.py
--- a/tests/pytorch/mock_sparse2/test_diag.py
+++ b/tests/pytorch/mock_sparse2/test_diag.py
@@ -44,7 +44,8 @@ def test_diag(val_shape, mat_shape):
     assert sp_mat.device == mat.device
     # row, col, val
     edge_index = torch.arange(len(val)).to(mat.device)
-    row, col, val = sp_mat.coo()
+    row, col = sp_mat.coo()
+    val = sp_mat.val
     assert torch.allclose(row, edge_index)
     assert torch.allclose(col, edge_index)
     assert torch.allclose(val, val)
diff --git a/tests/pytorch/mock_sparse2/test_elementwise_op_sp.py b/tests/pytorch/mock_sparse2/test_elementwise_op_sp.py
--- a/tests/pytorch/mock_sparse2/test_elementwise_op_sp.py
+++ b/tests/pytorch/mock_sparse2/test_elementwise_op_sp.py
@@ -12,7 +12,8 @@
     pytest.skip("skipping tests on win", allow_module_level=True)
 
 def all_close_sparse(A, row, col, val, shape):
-    rowA, colA, valA = A.coo()
+    rowA, colA = A.coo()
+    valA = A.val
     assert torch.allclose(rowA, row)
     assert torch.allclose(colA, col)
     assert torch.allclose(valA, val)
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
@@ -10,17 +10,12 @@
 if not sys.platform.startswith("linux"):
     pytest.skip("skipping tests on win", allow_module_level=True)
 
+
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("row", [[0, 0, 1, 2], (0, 1, 2, 4)])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
 @pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
-@pytest.mark.parametrize("shape", [None, (3, 5), (5, 3)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
 def test_create_from_coo(dense_dim, row, col, shape):
-    # Skip invalid matrices
-    if shape is not None and (
-        max(row) >= shape[0] or max(col) >= shape[1]
-    ):
-        return
-
     val_shape = (len(row),)
     if dense_dim is not None:
         val_shape += (dense_dim,)
@@ -33,7 +28,8 @@ def test_create_from_coo(dense_dim, row, col, shape):
     if shape is None:
         shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
 
-    mat_row, mat_col, mat_val = mat.coo()
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
 
     assert mat.shape == shape
     assert mat.nnz == row.numel()
@@ -44,7 +40,7 @@ def test_create_from_coo(dense_dim, row, col, shape):
 
 
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
 @pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
 @pytest.mark.parametrize("shape", [None, (3, 5)])
 def test_create_from_csr(dense_dim, indptr, indices, shape):
@@ -64,13 +60,15 @@ def test_create_from_csr(dense_dim, indptr, indices, shape):
     assert mat.shape == shape
     assert mat.nnz == indices.numel()
     assert mat.dtype == val.dtype
-    mat_indptr, mat_indices, mat_val = mat.csr()
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
     assert torch.allclose(mat_indptr, indptr)
     assert torch.allclose(mat_indices, indices)
     assert torch.allclose(mat_val, val)
 
+
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
 @pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
 @pytest.mark.parametrize("shape", [None, (5, 3)])
 def test_create_from_csc(dense_dim, indptr, indices, shape):
@@ -90,11 +88,13 @@ def test_create_from_csc(dense_dim, indptr, indices, shape):
     assert mat.shape == shape
     assert mat.nnz == indices.numel()
     assert mat.dtype == val.dtype
-    mat_indptr, mat_indices, mat_val = mat.csc()
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
     assert torch.allclose(mat_indptr, indptr)
     assert torch.allclose(mat_indices, indices)
     assert torch.allclose(mat_val, val)
 
+
 @pytest.mark.parametrize("val_shape", [(3), (3, 2)])
 def test_dense(val_shape):
     ctx = F.ctx()
@@ -110,6 +110,7 @@ def test_dense(val_shape):
     mat[row, col] = val
     assert torch.allclose(A_dense, mat)
 
+
 def test_set_val():
     ctx = F.ctx()
 
@@ -123,8 +124,9 @@ def test_set_val():
     A.val = new_val
     assert torch.allclose(new_val, A.val)
 
+
 @pytest.mark.parametrize("dense_dim", [None, 4])
-@pytest.mark.parametrize("indptr", [[0, 0, 1, 4], (0, 1, 2, 4)])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
 @pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
 @pytest.mark.parametrize("shape", [None, (3, 5)])
 def test_csr_to_coo(dense_dim, indptr, indices, shape):
@@ -140,9 +142,14 @@ def test_csr_to_coo(dense_dim, indptr, indices, shape):
     if shape is None:
         shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
 
-    row = torch.arange(0, indptr.shape[0] - 1).to(ctx).repeat_interleave(torch.diff(indptr))
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
     col = indices
-    mat_row, mat_col, mat_val = mat.coo()
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
 
     assert mat.shape == shape
     assert mat.nnz == row.numel()
@@ -152,3 +159,194 @@ def test_csr_to_coo(dense_dim, indptr, indices, shape):
     assert torch.allclose(mat_row, row)
     assert torch.allclose(mat_col, col)
 
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_coo(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = create_from_csc(indptr, indices, val, shape)
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    mat_row, mat_col = mat.coo()
+    mat_val = mat.val
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_row, row)
+    assert torch.allclose(mat_col, col)
+
+
+def _scatter_add(a, index, v=1):
+    index = index.tolist()
+    for i in index:
+        a[i] += v
+    return a
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csr(dense_dim, row, col, shape):
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = create_from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 4, 3, 2)])
+@pytest.mark.parametrize("shape", [None, (5, 3)])
+def test_csc_to_csr(dense_dim, indptr, indices, shape):
+    ctx = F.ctx()
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = create_from_csc(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csr()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (torch.max(indices).item() + 1, indptr.numel() - 1)
+
+    col = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+    row = indices
+    row, sort_index = row.sort(stable=True)
+    col = col[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[0] + 1).to(ctx)
+    indptr = _scatter_add(indptr, row + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = col
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("row", [(0, 0, 1, 2), (0, 1, 2, 4)])
+@pytest.mark.parametrize("col", [(0, 1, 2, 2), (1, 3, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (5, 5), (5, 6)])
+def test_coo_to_csc(dense_dim, row, col, shape):
+
+    val_shape = (len(row),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    row = torch.tensor(row).to(ctx)
+    col = torch.tensor(col).to(ctx)
+    mat = create_from_coo(row, col, val, shape)
+
+    if shape is None:
+        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
+
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)
+
+
+@pytest.mark.parametrize("dense_dim", [None, 4])
+@pytest.mark.parametrize("indptr", [(0, 0, 1, 4), (0, 1, 2, 4)])
+@pytest.mark.parametrize("indices", [(0, 1, 2, 3), (1, 2, 3, 4)])
+@pytest.mark.parametrize("shape", [None, (3, 5)])
+def test_csr_to_csc(dense_dim, indptr, indices, shape):
+    val_shape = (len(indices),)
+    if dense_dim is not None:
+        val_shape += (dense_dim,)
+    ctx = F.ctx()
+    val = torch.randn(val_shape).to(ctx)
+    indptr = torch.tensor(indptr).to(ctx)
+    indices = torch.tensor(indices).to(ctx)
+    mat = create_from_csr(indptr, indices, val, shape)
+    mat_indptr, mat_indices, value_indices = mat.csc()
+    mat_val = mat.val if value_indices is None else mat.val[value_indices]
+
+    if shape is None:
+        shape = (indptr.numel() - 1, torch.max(indices).item() + 1)
+
+    row = (
+        torch.arange(0, indptr.shape[0] - 1)
+        .to(ctx)
+        .repeat_interleave(torch.diff(indptr))
+    )
+
+    col = indices
+    col, sort_index = col.sort(stable=True)
+    row = row[sort_index]
+    val = val[sort_index]
+    indptr = torch.zeros(shape[1] + 1).to(ctx)
+    indptr = _scatter_add(indptr, col + 1)
+    indptr = torch.cumsum(indptr, 0).long()
+    indices = row
+
+    assert mat.shape == shape
+    assert mat.nnz == row.numel()
+    assert mat.device == row.device
+    assert mat.dtype == val.dtype
+    assert torch.allclose(mat_val, val)
+    assert torch.allclose(mat_indptr, indptr)
+    assert torch.allclose(mat_indices, indices)