[Sparse] Migrate DiagMatrix class (#4567)

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

* Update

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

docs/source/api/python/dgl.sparse_v0.rst

@@ -13,4 +13,24 @@ Sparse matrix class
 .. currentmodule:: dgl.mock_sparse
 
 .. autoclass:: SparseMatrix
-	:members:
\ No newline at end of file
+	:members: shape, nnz, dtype, device, row, col, val, __call__, indices, coo, csr, csc, dense
+
+.. autosummary::
+    :toctree: ../../generated/
+
+    create_from_coo
+    create_from_csr
+    create_from_csc
+
+Diagonal matrix class
+-------------------------
+.. currentmodule:: dgl.mock_sparse
+
+.. autoclass:: DiagMatrix
+	:members: val, shape, __call__, nnz, dtype, device, as_sparse
+
+.. autosummary::
+    :toctree: ../../generated/
+
+    diag
+    identity

docs/source/index.rst

@@ -55,6 +55,7 @@ Welcome to Deep Graph Library Tutorials and Documentation
    api/python/dgl.contrib.UnifiedTensor
    api/python/transforms
    api/python/udf
+   api/python/dgl.sparse_v0
 
 .. toctree::
    :maxdepth: 1

python/dgl/mock_sparse/__init__.py

 """dgl sparse class."""
+from .diag_matrix import *
 from .sp_matrix import *
 from .elementwise_op_sp import *

python/dgl/mock_sparse/diag_matrix.py

+"""dgl diagonal matrix module."""
+from typing import Optional, Tuple
+
+import torch
+
+from .sp_matrix import SparseMatrix, create_from_coo
+
+class DiagMatrix:
+    """Diagonal Matrix Class
+
+    Parameters
+    ----------
+    val : torch.Tensor
+        Diagonal of the matrix. It can take shape (N) or (N, D).
+    shape : tuple[int, int], optional
+        If not specified, it will be inferred from :attr:`val`, i.e.,
+        (N, N). Otherwise, :attr:`len(val)` must be equal to :attr:`min(shape)`.
+
+    Attributes
+    ----------
+    val : torch.Tensor
+        Diagonal of the matrix.
+    shape : tuple[int, int]
+        Shape of the matrix.
+    """
+    def __init__(self, val: torch.Tensor, shape: Optional[Tuple[int, int]] = None):
+        len_val = len(val)
+        if shape is not None:
+            assert len_val == min(shape), \
+            f'Expect len(val) to be min(shape), got {len_val} for len(val) and {shape} for shape.'
+        else:
+            shape = (len_val, len_val)
+        self.val = val
+        self.shape = shape
+
+    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
+
+        Returns
+        -------
+        int
+            The number of non-zero values in the matrix
+        """
+        return self.val.shape[0]
+
+    @property
+    def dtype(self) -> torch.dtype:
+        """Return the data type of the matrix
+
+        Returns
+        -------
+        torch.dtype
+            Data type of the matrix
+        """
+        return self.val.dtype
+
+    @property
+    def device(self) -> torch.device:
+        """Return the device of the matrix
+
+        Returns
+        -------
+        torch.device
+            Device of the matrix
+        """
+        return self.val.device
+
+    def as_sparse(self) -> SparseMatrix:
+        """Convert the diagonal matrix into a sparse matrix object
+
+        Returns
+        -------
+        SparseMatrix
+            The converted sparse matrix object
+
+        Example
+        -------
+
+        >>> import torch
+        >>> val = torch.ones(5)
+        >>> mat = diag(val)
+        >>> sp_mat = mat.as_sparse()
+        >>> print(sp_mat)
+        SparseMatrix(indices=tensor([[0, 1, 2, 3, 4],
+                                     [0, 1, 2, 3, 4]]),
+                     values=tensor([1., 1., 1., 1., 1.]),
+                     shape=(5, 5), nnz=5)
+        """
+        row = col = torch.arange(len(self.val)).to(self.device)
+        return create_from_coo(row=row, col=col, val=self.val, shape=self.shape)
+
+def diag(val: torch.Tensor, shape: Optional[Tuple[int, int]] = None) -> DiagMatrix:
+    """Create a diagonal matrix based on the diagonal values
+
+    Parameters
+    ----------
+    val : torch.Tensor
+        Diagonal of the matrix. It can take shape (N) or (N, D).
+    shape : tuple[int, int], optional
+        If not specified, it will be inferred from :attr:`val`, i.e.,
+        (N, N). Otherwise, :attr:`len(val)` must be equal to :attr:`min(shape)`.
+
+    Returns
+    -------
+    DiagMatrix
+        Diagonal matrix
+
+    Examples
+    --------
+
+    Case1: 5-by-5 diagonal matrix with scaler values on the diagonal
+
+    >>> import torch
+    >>> val = torch.ones(5)
+    >>> mat = diag(val)
+    >>> print(mat)
+    DiagMatrix(val=tensor([1., 1., 1., 1., 1.]),
+               shape=(5, 5))
+
+    Case2: 5-by-10 diagonal matrix with scaler values on the diagonal
+
+    >>> val = torch.ones(5)
+    >>> mat = diag(val, shape=(5, 10))
+    >>> print(mat)
+    DiagMatrix(val=tensor([1., 1., 1., 1., 1.]),
+               shape=(5, 10))
+
+    Case3: 5-by-5 diagonal matrix with tensor values on the diagonal
+
+    >>> val = torch.randn(5, 3)
+    >>> mat = diag(val)
+    >>> mat.shape
+    (5, 5)
+    >>> mat.nnz
+    5
+    """
+    # NOTE(Mufei): this may not be needed if DiagMatrix is simple enough
+    return DiagMatrix(val, shape)
+
+def identity(shape: Tuple[int, int],
+             d: Optional[int] = None,
+             dtype: Optional[torch.dtype] = None,
+             device: Optional[torch.device] = None) -> DiagMatrix:
+    """Create a diagonal matrix with ones on the diagonal and zeros elsewhere
+
+    Parameters
+    ----------
+    shape : tuple[int, int]
+        Shape of the matrix.
+    d : int, optional
+        If None, the diagonal entries will be scaler 1. Otherwise, the diagonal
+        entries will be a 1-valued tensor of shape (d).
+    dtype : torch.dtype, optional
+        The data type of the matrix
+    device : torch.device, optional
+        The device of the matrix
+
+    Returns
+    -------
+    DiagMatrix
+        Diagonal matrix
+
+    Examples
+    --------
+
+    Case1: 3-by-3 matrix with scaler diagonal values
+
+    [[1, 0, 0],
+     [0, 1, 0],
+     [0, 0, 1]]
+
+    >>> mat = identity(shape=(3, 3))
+    >>> print(mat)
+    DiagMatrix(val=tensor([1., 1., 1.]),
+               shape=(3, 3))
+
+    Case2: 3-by-5 matrix with scaler diagonal values
+
+    [[1, 0, 0, 0, 0],
+     [0, 1, 0, 0, 0],
+     [0, 0, 1, 0, 0]]
+
+    >>> mat = identity(shape=(3, 5))
+    >>> print(mat)
+    DiagMatrix(val=tensor([1., 1., 1.]),
+               shape=(3, 5))
+
+    Case3: 3-by-3 matrix with tensor diagonal values
+
+    >>> mat = identity(shape=(3, 3), d=2)
+    >>> mat.val
+    tensor([[1., 1.],
+            [1., 1.],
+            [1., 1.]])
+    >>> mat.shape
+    (3, 3)
+    >>> mat.nnz
+    3
+    """
+    len_val = min(shape)
+    if d is None:
+        val_shape = (len_val,)
+    else:
+        val_shape = (len_val, d)
+    val = torch.ones(val_shape, dtype=dtype, device=device)
+    return diag(val, shape)

python/dgl/mock_sparse/sp_matrix.py

@@ -60,10 +60,12 @@ class SparseMatrix:
             val = torch.ones(row.shape[0])
         self._val = val
         i = torch.cat((row.unsqueeze(0), col.unsqueeze(0)), 0)
-        if shape is not None:
-            self.adj = torch.sparse_coo_tensor(i, val, shape).coalesce()
-        else:
+        if shape is None:
             self.adj = torch.sparse_coo_tensor(i, val).coalesce()
+        else:
+            if len(val.shape) > 1:
+                shape += (val.shape[-1],)
+            self.adj = torch.sparse_coo_tensor(i, val, shape).coalesce()
 
     def __repr__(self):
         return f'SparseMatrix(indices={self.indices("COO")}, \nvalues={self.val}, \
@@ -236,18 +238,23 @@ class SparseMatrix:
 
 def create_from_coo(row: torch.Tensor,
                     col: torch.Tensor,
-                    val: Optional[torch.Tensor] = None) -> SparseMatrix:
+                    val: Optional[torch.Tensor] = None,
+                    shape: Optional[Tuple[int, int]] = None) -> SparseMatrix:
     """Create a sparse matrix from row and column coordinates.
 
     Parameters
     ----------
     row : tensor
-        The row indices of shape nnz.
+        The row indices of shape (nnz).
     col : tensor
-        The column indices of shape nnz.
+        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.
+    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.
 
     Returns
     -------
@@ -267,6 +274,13 @@ def create_from_coo(row: torch.Tensor,
                                  [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
+    SparseMatrix(indices=tensor([[1, 1, 2],
+                                 [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.
@@ -280,11 +294,12 @@ def create_from_coo(row: torch.Tensor,
                                 [3, 3]]),
                  shape=(3, 5), nnz=3)
     """
-    return SparseMatrix(row, col, val)
+    return SparseMatrix(row=row, col=col, val=val, shape=shape)
 
 def create_from_csr(indptr: torch.Tensor,
                     indices: torch.Tensor,
-                    val: Optional[torch.Tensor] = None) -> SparseMatrix:
+                    val: Optional[torch.Tensor] = None,
+                    shape: Optional[Tuple[int, int]] = None) -> SparseMatrix:
     """Create a sparse matrix from CSR indices.
 
     For row i of the sparse matrix
@@ -295,12 +310,16 @@ def create_from_csr(indptr: torch.Tensor,
     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.
+        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.
+    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.
 
     Returns
     -------
@@ -319,14 +338,15 @@ def create_from_csr(indptr: torch.Tensor,
     >>> indptr = torch.tensor([0, 1, 2, 5])
     >>> indices = torch.tensor([1, 2, 0, 1, 2])
     >>> A = create_from_csr(indptr, indices)
+    >>> 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)
+    >>> # Specify shape
+    >>> A = create_from_csr(indptr, indices, shape=(5, 3))
     >>> A.shape
-    (3, 3)
-    >>> A.row
-    tensor([0, 1, 2, 2, 2])
-    >>> A.val
-    tensor([1., 1., 1., 1., 1.])
-    >>> A.nnz
-    5
+    (5, 3)
 
     Case2: Sparse matrix with scalar/vector values. Following example is with
     vector data.
@@ -344,11 +364,12 @@ def create_from_csr(indptr: torch.Tensor,
     adj_coo = adj_csr.to_sparse_coo().coalesce()
     row, col = adj_coo.indices()
 
-    return SparseMatrix(row, col, val)
+    return SparseMatrix(row=row, col=col, val=val, shape=shape)
 
 def create_from_csc(indptr: torch.Tensor,
                     indices: torch.Tensor,
-                    val: Optional[torch.Tensor] = None) -> SparseMatrix:
+                    val: Optional[torch.Tensor] = None,
+                    shape: Optional[Tuple[int, int]] = None) -> SparseMatrix:
     """Create a sparse matrix from CSC indices.
 
     For column i of the sparse matrix
@@ -365,6 +386,10 @@ def create_from_csc(indptr: torch.Tensor,
     val : tensor, optional
         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.
 
     Returns
     -------
@@ -383,14 +408,15 @@ def create_from_csc(indptr: torch.Tensor,
     >>> indptr = torch.tensor([0, 1, 3, 5])
     >>> indices = torch.tensor([2, 0, 2, 1, 2])
     >>> A = create_from_csc(indptr, indices)
+    >>> 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)
+    >>> # Specify shape
+    >>> A = create_from_csc(indptr, indices, shape=(5, 3))
     >>> A.shape
-    (3, 3)
-    >>> A.row
-    tensor([0, 1, 2, 2, 2])
-    >>> A.val
-    tensor([1., 1., 1., 1., 1.])
-    >>> A.nnz
-    5
+    (5, 3)
 
     Case2: Sparse matrix with scalar/vector values. Following example is with
     vector data.
@@ -408,4 +434,4 @@ def create_from_csc(indptr: torch.Tensor,
     adj_coo = adj_csr.to_sparse_coo().coalesce()
     col, row = adj_coo.indices()
 
-    return SparseMatrix(row, col, val)
+    return SparseMatrix(row=row, col=col, val=val, shape=shape)

tests/pytorch/mock_sparse/test_diag.py

+import pytest
+import torch
+
+import backend as F
+
+from dgl.mock_sparse import diag, identity, DiagMatrix
+
+@pytest.mark.parametrize('val_shape', [(3,), (3, 2)])
+@pytest.mark.parametrize('mat_shape', [None, (3, 5), (5, 3)])
+def test_diag(val_shape, mat_shape):
+    # creation
+    val = torch.randn(val_shape).to(F.ctx())
+    mat = diag(val, mat_shape)
+
+    # val, shape attributes
+    assert torch.allclose(mat.val, val)
+    if mat_shape is None:
+        mat_shape = (val_shape[0], val_shape[0])
+    assert mat.shape == mat_shape
+
+    # __call__
+    val = torch.randn(val_shape).to(F.ctx())
+    mat = mat(val)
+    assert torch.allclose(mat.val, val)
+
+    # nnz
+    assert mat.nnz == val.shape[0]
+    # dtype
+    assert mat.dtype == val.dtype
+    # device
+    assert mat.device == val.device
+
+    # as_sparse
+    sp_mat = mat.as_sparse()
+    # shape
+    assert sp_mat.shape == mat_shape
+    # nnz
+    assert sp_mat.nnz == mat.nnz
+    # dtype
+    assert sp_mat.dtype == mat.dtype
+    # device
+    assert sp_mat.device == mat.device
+    # row, col, val
+    edge_index = torch.arange(len(val)).to(mat.device)
+    assert torch.allclose(sp_mat.row, edge_index)
+    assert torch.allclose(sp_mat.col, edge_index)
+    assert torch.allclose(sp_mat.val, val)
+
+@pytest.mark.parametrize('shape', [(3, 3), (3, 5), (5, 3)])
+@pytest.mark.parametrize('d', [None, 2])
+def test_identity(shape, d):
+    # creation
+    mat = identity(shape, d, device=F.ctx())
+    # type
+    assert isinstance(mat, DiagMatrix)
+    # shape
+    assert mat.shape == shape
+    # val
+    len_val = min(shape)
+    if d is None:
+        val_shape = (len_val)
+    else:
+        val_shape = (len_val, d)
+    val = torch.ones(val_shape, device=F.ctx())
+    assert torch.allclose(val, mat.val)