[Sparse] Set Value and Dense (#4897)

* Update

* dense

* Update

* Update

* Update

* Update

* Update
Co-authored-by: Ubuntu <ubuntu@ip-172-31-36-188.ap-northeast-1.compute.internal>

dgl_sparse/include/sparse/sparse_matrix.h

@@ -129,6 +129,12 @@ class SparseMatrix : public torch::CustomClassHolder {
   /** @return {row, col, value} tensors in the CSC format. */
   std::vector<torch::Tensor> CSCTensors();
 
+  /**
+   * @brief Set non-zero values of the sparse matrix
+   * @param value Values of the sparse matrix
+   */
+  void SetValue(torch::Tensor value);
+
  private:
   /** @brief Create the COO format for the sparse matrix internally */
   void _CreateCOO();

dgl_sparse/src/python_binding.cc

@@ -19,7 +19,8 @@ TORCH_LIBRARY(dgl_sparse, m) {
       .def("shape", &SparseMatrix::shape)
       .def("coo", &SparseMatrix::COOTensors)
       .def("csr", &SparseMatrix::CSRTensors)
-      .def("csc", &SparseMatrix::CSCTensors);
+      .def("csc", &SparseMatrix::CSCTensors)
+      .def("set_val", &SparseMatrix::SetValue);
   m.def("create_from_coo", &CreateFromCOO)
       .def("create_from_csr", &CreateFromCSR)
       .def("create_from_csc", &CreateFromCSC)

dgl_sparse/src/sparse_matrix.cc

@@ -112,6 +112,8 @@ std::vector<torch::Tensor> SparseMatrix::CSCTensors() {
   return {csc->indptr, csc->indices, val};
 }
 
+void SparseMatrix::SetValue(torch::Tensor value) { value_ = value; }
+
 // TODO(zhenkun): format conversion
 void SparseMatrix::_CreateCOO() {}
 void SparseMatrix::_CreateCSR() {}

python/dgl/mock_sparse2/diag_matrix.py

@@ -42,37 +42,6 @@ class DiagMatrix:
     def __repr__(self):
         return f"DiagMatrix(val={self.val}, \nshape={self.shape})"
 
-    def __call__(self, x: torch.Tensor):
-        """Create a new diagonal matrix with the same shape as self
-        but different values.
-
-        Parameters
-        ----------
-        x : torch.Tensor
-            Values of the diagonal matrix
-
-        Returns
-        -------
-        DiagMatrix
-            Diagonal matrix
-
-        Examples
-        --------
-
-        >>> import torch
-        >>> val = torch.ones(5)
-        >>> mat = diag(val)
-        >>> print(mat)
-        DiagMatrix(val=tensor([1., 1., 1., 1., 1.]),
-                   shape=(5, 5))
-        >>> val = torch.ones(5) + 1
-        >>> mat = mat(val)
-        >>> print(mat)
-        DiagMatrix(val=tensor([2., 2., 2., 2., 2.]),
-                   shape=(5, 5))
-        """
-        return diag(x, self.shape)
-
     @property
     def nnz(self) -> int:
         """Return the number of non-zero values in the matrix

python/dgl/mock_sparse2/sparse_matrix.py

@@ -21,6 +21,17 @@ class SparseMatrix:
         """
         return self.c_sparse_matrix.val()
 
+    @val.setter
+    def val(self, x: torch.Tensor):
+        """Set the non-zero values inplace.
+
+        Parameters
+        ----------
+        x : torch.Tensor, optional
+            The values of shape (nnz) or (nnz, D)
+        """
+        self.c_sparse_matrix.set_val(x)
+
     @property
     def shape(self) -> Tuple[int]:
         """Shape of the sparse matrix.
@@ -90,8 +101,8 @@ class SparseMatrix:
             raise NotImplementedError
 
     def __repr__(self):
-        return f'SparseMatrix(indices={self.indices("COO")}, \nvalues={self.val}, \
-                \nshape={self.shape}, nnz={self.nnz})'
+        return f'SparseMatrix(indices={self.indices("COO")}, \
+                \nvalues={self.val}, \nshape={self.shape}, nnz={self.nnz})'
 
     def coo(self) -> Tuple[torch.Tensor, ...]:
         """Get the coordinate (COO) representation of the sparse matrix.
@@ -123,6 +134,20 @@ class SparseMatrix:
         """
         return self.c_sparse_matrix.csc()
 
+    def dense(self) -> torch.Tensor:
+        """Return a dense representation of the matrix.
+
+        Returns
+        -------
+        torch.Tensor
+            Dense representation of the sparse matrix.
+        """
+        row, col, val = self.coo()
+        shape = self.shape + val.shape[1:]
+        mat = torch.zeros(shape, device=self.device)
+        mat[row, col] = val
+        return mat
+
 
 def create_from_coo(
     row: torch.Tensor,
@@ -139,12 +164,12 @@ def create_from_coo(
     col : tensor
         The column indices of shape (nnz).
     val : tensor, optional
-        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of shape (nnz)
-        filled by 1.
+        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of
+        shape (nnz) filled by 1.
     shape : tuple[int, int], optional
-        If not specified, it will be inferred from :attr:`row` and :attr:`col`, i.e.,
-        (row.max() + 1, col.max() + 1). Otherwise, :attr:`shape` should be no smaller
-        than this.
+        If not specified, it will be inferred from :attr:`row` and :attr:`col`,
+        i.e., (row.max() + 1, col.max() + 1). Otherwise, :attr:`shape` should
+        be no smaller than this.
 
     Returns
     -------
@@ -156,16 +181,16 @@ def create_from_coo(
 
     Case1: Sparse matrix with row and column indices without values.
 
-    >>> src = torch.tensor([1, 1, 2])
-    >>> dst = torch.tensor([2, 4, 3])
-    >>> A = create_from_coo(src, dst)
+    >>> dst = torch.tensor([1, 1, 2])
+    >>> src = torch.tensor([2, 4, 3])
+    >>> 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)
     >>> # Specify shape
-    >>> A = create_from_coo(src, dst, shape=(5, 5))
+    >>> A = create_from_coo(dst, src, shape=(5, 5))
     >>> print(A)
     SparseMatrix(indices=tensor([[1, 1, 2],
                                  [2, 4, 3]]),
@@ -175,8 +200,8 @@ def create_from_coo(
     Case2: Sparse matrix with scalar/vector values. Following example is with
     vector data.
 
-    >>> val = torch.tensor([[1, 1], [2, 2], [3, 3]])
-    >>> A = create_from_coo(src, dst, val)
+    >>> 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],
@@ -187,7 +212,7 @@ def create_from_coo(
     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])
+        val = torch.ones(row.shape[0]).to(row.device)
 
     return SparseMatrix(
         torch.ops.dgl_sparse.create_from_coo(row, col, val, shape)
@@ -206,22 +231,24 @@ def create_from_csr(
 
     For row i of the sparse matrix
 
-    - the column indices of the nonzero entries are stored in ``indices[indptr[i]: indptr[i+1]]``
+    - the column indices of the nonzero entries are stored in
+      ``indices[indptr[i]: indptr[i+1]]``
     - the corresponding values are stored in ``val[indptr[i]: indptr[i+1]]``
 
     Parameters
     ----------
     indptr : tensor
-        Pointer to the column indices of shape (N + 1), where N is the number of rows.
+        Pointer to the column indices of shape (N + 1), where N is the number
+        of rows.
     indices : tensor
         The column indices of shape (nnz).
     val : tensor, optional
-        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of shape (nnz)
-        filled by 1.
+        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of
+        shape (nnz) filled by 1.
     shape : tuple[int, int], optional
-        If not specified, it will be inferred from :attr:`indptr` and :attr:`indices`, i.e.,
-        (len(indptr) - 1, indices.max() + 1). Otherwise, :attr:`shape` should be no smaller
-        than this.
+        If not specified, it will be inferred from :attr:`indptr` and
+        :attr:`indices`, i.e., (len(indptr) - 1, indices.max() + 1). Otherwise,
+        :attr:`shape` should be no smaller than this.
 
     Returns
     -------
@@ -271,7 +298,7 @@ def create_from_csr(
     if shape is None:
         shape = (indptr.shape[0] - 1, torch.max(indices) + 1)
     if val is None:
-        val = torch.ones(indices.shape[0])
+        val = torch.ones(indices.shape[0]).to(indptr.device)
 
     return SparseMatrix(
         torch.ops.dgl_sparse.create_from_csr(indptr, indices, val, shape)
@@ -290,22 +317,24 @@ def create_from_csc(
 
     For column i of the sparse matrix
 
-    - the row indices of the nonzero entries are stored in ``indices[indptr[i]: indptr[i+1]]``
+    - the row indices of the nonzero entries are stored in
+      ``indices[indptr[i]: indptr[i+1]]``
     - the corresponding values are stored in ``val[indptr[i]: indptr[i+1]]``
 
     Parameters
     ----------
     indptr : tensor
-        Pointer to the row indices of shape N + 1, where N is the number of columns.
+        Pointer to the row indices of shape N + 1, where N is the
+        number of columns.
     indices : tensor
         The row indices of shape nnz.
     val : tensor, optional
-        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of shape (nnz)
-        filled by 1.
+        The values of shape (nnz) or (nnz, D). If None, it will be a tensor of
+        shape (nnz) filled by 1.
     shape : tuple[int, int], optional
-        If not specified, it will be inferred from :attr:`indptr` and :attr:`indices`, i.e.,
-        (indices.max() + 1, len(indptr) - 1). Otherwise, :attr:`shape` should be no smaller
-        than this.
+        If not specified, it will be inferred from :attr:`indptr` and
+        :attr:`indices`, i.e., (indices.max() + 1, len(indptr) - 1). Otherwise,
+        :attr:`shape` should be no smaller than this.
 
     Returns
     -------
@@ -355,7 +384,7 @@ def create_from_csc(
     if shape is None:
         shape = (torch.max(indices) + 1, indptr.shape[0] - 1)
     if val is None:
-        val = torch.ones(indices.shape[0])
+        val = torch.ones(indices.shape[0]).to(indptr.device)
 
     return SparseMatrix(
         torch.ops.dgl_sparse.create_from_csc(indptr, indices, val, shape)

tests/pytorch/mock_sparse2/test_diag.py

@@ -24,8 +24,6 @@ def test_diag(val_shape, mat_shape):
 
     # __call__
     val = torch.randn(val_shape)
-    mat = mat(val)
-    assert torch.allclose(mat.val, val)
 
     # nnz
     assert mat.nnz == val.shape[0]

tests/pytorch/mock_sparse2/test_sparse_matrix.py

@@ -2,6 +2,8 @@ import pytest
 import torch
 import sys
 
+import backend as F
+
 from dgl.mock_sparse2 import create_from_coo, create_from_csr, create_from_csc
 
 # FIXME: Skipping tests on win.
@@ -22,14 +24,15 @@ def test_create_from_coo(dense_dim, row, col, shape):
     val_shape = (len(row),)
     if dense_dim is not None:
         val_shape += (dense_dim,)
-    val = torch.randn(val_shape)
-    row = torch.tensor(row)
-    col = torch.tensor(col)
+    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_row, mat_col, mat_val = mat.coo()
 
     assert mat.shape == shape
@@ -48,9 +51,10 @@ def test_create_from_csr(dense_dim, indptr, indices, shape):
     val_shape = (len(indices),)
     if dense_dim is not None:
         val_shape += (dense_dim,)
-    val = torch.randn(val_shape)
-    indptr = torch.tensor(indptr)
-    indices = torch.tensor(indices)
+    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)
 
     if shape is None:
@@ -73,9 +77,10 @@ def test_create_from_csc(dense_dim, indptr, indices, shape):
     val_shape = (len(indices),)
     if dense_dim is not None:
         val_shape += (dense_dim,)
-    val = torch.randn(val_shape)
-    indptr = torch.tensor(indptr)
-    indices = torch.tensor(indices)
+    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_csc(indptr, indices, val, shape)
 
     if shape is None:
@@ -89,3 +94,31 @@ def test_create_from_csc(dense_dim, indptr, indices, shape):
     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()
+
+    row = torch.tensor([1, 1, 2]).to(ctx)
+    col = torch.tensor([2, 4, 3]).to(ctx)
+    val = torch.randn(val_shape).to(ctx)
+    A = create_from_coo(row, col, val)
+    A_dense = A.dense()
+
+    shape = A.shape + val.shape[1:]
+    mat = torch.zeros(shape, device=ctx)
+    mat[row, col] = val
+    assert torch.allclose(A_dense, mat)
+
+def test_set_val():
+    ctx = F.ctx()
+
+    row = torch.tensor([1, 1, 2]).to(ctx)
+    col = torch.tensor([2, 4, 3]).to(ctx)
+    nnz = len(row)
+    old_val = torch.ones(nnz).to(ctx)
+    A = create_from_coo(row, col, old_val)
+
+    new_val = torch.zeros(nnz).to(ctx)
+    A.val = new_val
+    assert torch.allclose(new_val, A.val)