Repository: dmlc/dgl

Issue:
[Sparse] Convert Torch sparse tensor to SparseMatrix without overhead.
## 🔨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/sparse/sparse_matrix.py b/python/dgl/sparse/sparse_matrix.py
--- a/python/dgl/sparse/sparse_matrix.py
+++ b/python/dgl/sparse/sparse_matrix.py
@@ -109,12 +109,35 @@ def coo(self) -> Tuple[torch.Tensor, torch.Tensor]:
         --------
 
         >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]])
-        >>> A = from_coo(dst, src)
+        >>> A = dglsp.spmatrix(indices)
         >>> A.coo()
         (tensor([1, 2, 1]), tensor([2, 4, 3]))
         """
         return self.c_sparse_matrix.coo()
 
+    def indices(self) -> torch.Tensor:
+        r"""Returns the coordinate list (COO) representation in one tensor with
+        shape ``(2, nnz)``.
+
+        See `COO in Wikipedia <https://en.wikipedia.org/wiki/
+        Sparse_matrix#Coordinate_list_(COO)>`_.
+
+        Returns
+        -------
+        torch.Tensor
+            Stacked COO tensor with shape ``(2, nnz)``.
+
+        Examples
+        --------
+
+        >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]])
+        >>> A = dglsp.spmatrix(indices)
+        >>> A.indices()
+        tensor([[1, 2, 1],
+                [2, 4, 3]])
+        """
+        return self.c_sparse_matrix.indices()
+
     def csr(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         r"""Returns the compressed sparse row (CSR) representation of the sparse
         matrix.
@@ -140,7 +163,7 @@ def csr(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         --------
 
         >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]])
-        >>> A = from_coo(dst, src)
+        >>> A = dglsp.spmatrix(indices)
         >>> A.csr()
         (tensor([0, 0, 2, 3]), tensor([2, 3, 4]), tensor([0, 2, 1]))
         """
@@ -171,7 +194,7 @@ def csc(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         --------
 
         >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]])
-        >>> A = from_coo(dst, src)
+        >>> A = dglsp.spmatrix(indices)
         >>> A.csc()
         (tensor([0, 0, 0, 1, 2, 3]), tensor([1, 1, 2]), tensor([0, 2, 1]))
         """
@@ -521,7 +544,18 @@ def spmatrix(
                                 [3., 3.]]),
                  shape=(3, 5), nnz=3, val_size=(2,))
     """
-    return from_coo(indices[0], indices[1], val, shape)
+    if shape is None:
+        shape = (
+            torch.max(indices[0]).item() + 1,
+            torch.max(indices[1]).item() + 1,
+        )
+    if val is None:
+        val = torch.ones(indices.shape[1]).to(indices.device)
+
+    assert (
+        val.dim() <= 2
+    ), "The values of a SparseMatrix can only be scalars or vectors."
+    return SparseMatrix(torch.ops.dgl_sparse.from_coo(indices, val, shape))
 
 
 def from_coo(
@@ -599,16 +633,8 @@ def from_coo(
                                 [3., 3.]]),
                  shape=(3, 5), nnz=3, val_size=(2,))
     """
-    if shape is None:
-        shape = (torch.max(row).item() + 1, torch.max(col).item() + 1)
-    if val is None:
-        val = torch.ones(row.shape[0]).to(row.device)
-
-    assert (
-        val.dim() <= 2
-    ), "The values of a SparseMatrix can only be scalars or vectors."
-
-    return SparseMatrix(torch.ops.dgl_sparse.from_coo(row, col, val, shape))
+    assert row.shape[0] == col.shape[0]
+    return spmatrix(torch.stack([row, col]), val, shape)
 
 
 def from_csr(
@@ -833,6 +859,171 @@ def val_like(mat: SparseMatrix, val: torch.Tensor) -> SparseMatrix:
     return SparseMatrix(torch.ops.dgl_sparse.val_like(mat.c_sparse_matrix, val))
 
 
+def from_torch_sparse(torch_sparse_tensor: torch.Tensor) -> SparseMatrix:
+    """Creates a sparse matrix from a torch sparse tensor, which can have coo,
+    csr, or csc layout.
+
+    Parameters
+    ----------
+    torch_sparse_tensor : torch.Tensor
+        Torch sparse tensor
+
+    Returns
+    -------
+    SparseMatrix
+        Sparse matrix
+
+    Examples
+    --------
+
+    >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]])
+    >>> val = torch.ones(3)
+    >>> torch_coo = torch.sparse_coo_tensor(indices, val)
+    >>> dglsp.from_torch_sparse(torch_coo)
+    SparseMatrix(indices=tensor([[1, 1, 2],
+                                 [2, 4, 3]]),
+                 values=tensor([1., 1., 1.]),
+                 shape=(3, 5), nnz=3)
+    """
+    assert torch_sparse_tensor.layout in (
+        torch.sparse_coo,
+        torch.sparse_csr,
+        torch.sparse_csc,
+    ), (
+        f"Cannot convert Pytorch sparse tensor with layout "
+        f"{torch_sparse_tensor.layout} to DGL sparse."
+    )
+    if torch_sparse_tensor.layout == torch.sparse_coo:
+        # Use ._indices() and ._values() to access uncoalesced indices and
+        # values.
+        return spmatrix(
+            torch_sparse_tensor._indices(),
+            torch_sparse_tensor._values(),
+            torch_sparse_tensor.shape[:2],
+        )
+    elif torch_sparse_tensor.layout == torch.sparse_csr:
+        return from_csr(
+            torch_sparse_tensor.crow_indices(),
+            torch_sparse_tensor.col_indices(),
+            torch_sparse_tensor.values(),
+            torch_sparse_tensor.shape[:2],
+        )
+    else:
+        return from_csc(
+            torch_sparse_tensor.ccol_indices(),
+            torch_sparse_tensor.row_indices(),
+            torch_sparse_tensor.values(),
+            torch_sparse_tensor.shape[:2],
+        )
+
+
+def to_torch_sparse_coo(spmat: SparseMatrix) -> torch.Tensor:
+    """Creates a torch sparse coo tensor from a sparse matrix.
+
+    Parameters
+    ----------
+    spmat : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    torch.Tensor
+        torch tensor with torch.sparse_coo layout
+
+    Examples
+    --------
+
+    >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]])
+    >>> val = torch.ones(3)
+    >>> spmat = dglsp.spmatrix(indices, val)
+    >>> dglsp.to_torch_sparse_coo(spmat)
+    tensor(indices=tensor([[1, 1, 2],
+                           [2, 4, 3]]),
+           values=tensor([1., 1., 1.]),
+           size=(3, 5), nnz=3, layout=torch.sparse_coo)
+    """
+    shape = spmat.shape
+    if spmat.val.dim() > 1:
+        shape += spmat.val.shape[1:]
+    return torch.sparse_coo_tensor(spmat.indices(), spmat.val, shape)
+
+
+def to_torch_sparse_csr(spmat: SparseMatrix) -> torch.Tensor:
+    """Creates a torch sparse csr tensor from a sparse matrix.
+
+    Note that converting a sparse matrix to torch csr tensor could change the
+    order of non-zero values.
+
+    Parameters
+    ----------
+    spmat : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    torch.Tensor
+        Torch tensor with torch.sparse_csr layout
+
+    Examples
+    --------
+
+    >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]])
+    >>> val = torch.arange(3)
+    >>> spmat = dglsp.spmatrix(indices, val)
+    >>> dglsp.to_torch_sparse_csr(spmat)
+    tensor(crow_indices=tensor([0, 0, 2, 3]),
+           col_indices=tensor([2, 3, 4]),
+           values=tensor([0, 2, 1]), size=(3, 5), nnz=3,
+           layout=torch.sparse_csr)
+    """
+    shape = spmat.shape
+    if spmat.val.dim() > 1:
+        shape += spmat.val.shape[1:]
+    indptr, indices, value_indices = spmat.csr()
+    val = spmat.val
+    if value_indices is not None:
+        val = val[value_indices]
+    return torch.sparse_csr_tensor(indptr, indices, val, shape)
+
+
+def to_torch_sparse_csc(spmat: SparseMatrix) -> torch.Tensor:
+    """Creates a torch sparse csc tensor from a sparse matrix.
+
+    Note that converting a sparse matrix to torch csc tensor could change the
+    order of non-zero values.
+
+    Parameters
+    ----------
+    spmat : SparseMatrix
+        Sparse matrix
+
+    Returns
+    -------
+    torch.Tensor
+        Torch tensor with torch.sparse_csc layout
+
+    Examples
+    --------
+
+    >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]])
+    >>> val = torch.arange(3)
+    >>> spmat = dglsp.spmatrix(indices, val)
+    >>> dglsp.to_torch_sparse_csc(spmat)
+    tensor(ccol_indices=tensor([0, 0, 0, 1, 2, 3]),
+           row_indices=tensor([1, 1, 2]),
+           values=tensor([0, 2, 1]), size=(3, 5), nnz=3,
+           layout=torch.sparse_csc)
+    """
+    shape = spmat.shape
+    if spmat.val.dim() > 1:
+        shape += spmat.val.shape[1:]
+    indptr, indices, value_indices = spmat.csc()
+    val = spmat.val
+    if value_indices is not None:
+        val = val[value_indices]
+    return torch.sparse_csc_tensor(indptr, indices, val, shape)
+
+
 def _sparse_matrix_str(spmat: SparseMatrix) -> str:
     """Internal function for converting a sparse matrix to string
     representation.


Test patch:
diff --git a/tests/python/pytorch/sparse/test_sparse_matrix.py b/tests/python/pytorch/sparse/test_sparse_matrix.py
--- a/tests/python/pytorch/sparse/test_sparse_matrix.py
+++ b/tests/python/pytorch/sparse/test_sparse_matrix.py
@@ -5,7 +5,16 @@
 import pytest
 import torch
 
-from dgl.sparse import from_coo, from_csc, from_csr, val_like
+from dgl.sparse import (
+    from_coo,
+    from_csc,
+    from_csr,
+    from_torch_sparse,
+    to_torch_sparse_coo,
+    to_torch_sparse_csc,
+    to_torch_sparse_csr,
+    val_like,
+)
 
 
 @pytest.mark.parametrize("dense_dim", [None, 4])
@@ -502,3 +511,98 @@ def test_sparse_matrix_transpose(dense_dim, row, col, extra_shape):
     assert torch.allclose(mat_val, val)
     assert torch.allclose(mat_row, col)
     assert torch.allclose(mat_col, row)
+
+
+@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("nz_dim", [None, 2])
+@pytest.mark.parametrize("shape", [(5, 5), (6, 7)])
+def test_torch_sparse_coo_conversion(row, col, nz_dim, shape):
+    dev = F.ctx()
+    row = torch.tensor(row).to(dev)
+    col = torch.tensor(col).to(dev)
+    indices = torch.stack([row, col])
+    torch_sparse_shape = shape
+    val_shape = (row.shape[0],)
+    if nz_dim is not None:
+        torch_sparse_shape += (nz_dim,)
+        val_shape += (nz_dim,)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_coo = torch.sparse_coo_tensor(indices, val, torch_sparse_shape)
+    spmat = from_torch_sparse(torch_sparse_coo)
+
+    def _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo):
+        assert torch_sparse_coo.layout == torch.sparse_coo
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert (
+            spmat.indices().data_ptr() == torch_sparse_coo._indices().data_ptr()
+        )
+        assert spmat.val.data_ptr() == torch_sparse_coo._values().data_ptr()
+        assert spmat.shape == torch_sparse_coo.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo)
+    torch_sparse_coo = to_torch_sparse_coo(spmat)
+    _assert_spmat_equal_to_torch_sparse_coo(spmat, torch_sparse_coo)
+
+
+@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", [(3, 5), (3, 7)])
+def test_torch_sparse_csr_conversion(indptr, indices, shape):
+    dev = F.ctx()
+    indptr = torch.tensor(indptr).to(dev)
+    indices = torch.tensor(indices).to(dev)
+    torch_sparse_shape = shape
+    val_shape = (indices.shape[0],)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_csr = torch.sparse_csr_tensor(
+        indptr, indices, val, torch_sparse_shape
+    )
+    spmat = from_torch_sparse(torch_sparse_csr)
+
+    def _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr):
+        indptr, indices, value_indices = spmat.csr()
+        assert torch_sparse_csr.layout == torch.sparse_csr
+        assert value_indices is None
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert indptr.data_ptr() == torch_sparse_csr.crow_indices().data_ptr()
+        assert indices.data_ptr() == torch_sparse_csr.col_indices().data_ptr()
+        assert spmat.val.data_ptr() == torch_sparse_csr.values().data_ptr()
+        assert spmat.shape == torch_sparse_csr.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr)
+    torch_sparse_csr = to_torch_sparse_csr(spmat)
+    _assert_spmat_equal_to_torch_sparse_csr(spmat, torch_sparse_csr)
+
+
+@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", [(8, 3), (5, 3)])
+def test_torch_sparse_csc_conversion(indptr, indices, shape):
+    dev = F.ctx()
+    indptr = torch.tensor(indptr).to(dev)
+    indices = torch.tensor(indices).to(dev)
+    torch_sparse_shape = shape
+    val_shape = (indices.shape[0],)
+    val = torch.randn(val_shape).to(dev)
+    torch_sparse_csc = torch.sparse_csc_tensor(
+        indptr, indices, val, torch_sparse_shape
+    )
+    spmat = from_torch_sparse(torch_sparse_csc)
+
+    def _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc):
+        indptr, indices, value_indices = spmat.csc()
+        assert torch_sparse_csc.layout == torch.sparse_csc
+        assert value_indices is None
+        # Use .data_ptr() to check whether indices and values are on the same
+        # memory address
+        assert indptr.data_ptr() == torch_sparse_csc.ccol_indices().data_ptr()
+        assert indices.data_ptr() == torch_sparse_csc.row_indices().data_ptr()
+        assert spmat.val.data_ptr() == torch_sparse_csc.values().data_ptr()
+        assert spmat.shape == torch_sparse_csc.shape[:2]
+
+    _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc)
+    torch_sparse_csc = to_torch_sparse_csc(spmat)
+    _assert_spmat_equal_to_torch_sparse_csc(spmat, torch_sparse_csc)