Skip to content

GitLab

Unverified Commit fa74ea70 authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Sparse] Polish diag_matrix.py and sparse_matrix.py. (#5178) 



* polish diagonal

* polsih

* more

* more

* add paramaters-back
Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent 0868f179
Show whitespace changes
Inline Side-by-side
Showing with 98 additions and 109 deletions
python/dgl/sparse/diag_matrix.py
View file @ fa74ea70
......@@ -8,22 +8,15 @@ from .sparse_matrix import from_coo, SparseMatrix
class DiagMatrix:
"""Diagonal Matrix Class
r"""Class for diagonal matrix.
Parameters
----------
val : torch.Tensor
Diagonal of the matrix. It can take shape (N) or (N, D).
Diagonal of the matrix, in 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.
If specified, :attr:`len(val)` must be equal to :attr:`min(shape)`,
otherwise, it will be inferred from :attr:`val`, i.e., ``(N, N)``
"""
def __init__(
......@@ -40,71 +33,71 @@ class DiagMatrix:
self._val = val
self._shape = shape
def __repr__(self):
return _diag_matrix_str(self)
@property
def val(self) -> torch.Tensor:
"""Get the values of the nonzero elements.
"""Returns the values of the non-zero elements.
Returns
-------
torch.Tensor
Values of the nonzero elements
Values of the non-zero elements
"""
return self._val
@property
def shape(self) -> Tuple[int]:
"""Shape of the sparse matrix.
"""Returns the shape of the diagonal matrix.
Returns
-------
Tuple[int]
The shape of the matrix
The shape of the diagonal matrix
"""
return self._shape
def __repr__(self):
return _diag_matrix_str(self)
@property
def nnz(self) -> int:
"""Return the number of non-zero values in the matrix
"""Returns the number of non-zero elements in the diagonal matrix.
Returns
-------
int
The number of non-zero values in the matrix
The number of non-zero elements in the diagonal matrix
"""
return self.val.shape[0]
@property
def dtype(self) -> torch.dtype:
"""Return the data type of the matrix
"""Returns the data type of the diagonal matrix.
Returns
-------
torch.dtype
Data type of the matrix
Data type of the diagonal matrix
"""
return self.val.dtype
@property
def device(self) -> torch.device:
"""Return the device of the matrix
"""Returns the device the diagonal matrix is on.
Returns
-------
torch.device
Device of the matrix
The device the diagonal matrix is on
"""
return self.val.device
def to_sparse(self) -> SparseMatrix:
"""Convert the diagonal matrix into a sparse matrix object
"""Returns a copy in sparse matrix format of the diagonal matrix.
Returns
-------
SparseMatrix
The converted sparse matrix object
The copy in sparse matrix format
Example
-------
......@@ -123,12 +116,12 @@ class DiagMatrix:
return from_coo(row=row, col=col, val=self.val, shape=self.shape)
def to_dense(self) -> torch.Tensor:
"""Return a dense representation of the matrix.
"""Returns a copy in dense matrix format of the diagonal matrix.
Returns
-------
torch.Tensor
Dense representation of the diagonal matrix.
The copy in dense matrix format
"""
val = self.val
device = self.device
......@@ -148,12 +141,12 @@ class DiagMatrix:
return self.transpose()
def transpose(self):
"""Return the transpose of the matrix.
"""Returns a matrix that is a transposed version of the diagonal matrix.
Returns
-------
DiagMatrix
The transpose of the matrix.
The transpose of the matrix
Example
--------
......@@ -168,22 +161,22 @@ class DiagMatrix:
return DiagMatrix(self.val, self.shape[::-1])
def to(self, device=None, dtype=None):
"""Perform matrix dtype and/or device conversion. If the target device
"""Performs matrix dtype and/or device conversion. If the target device
and dtype are already in use, the original matrix will be returned.
Parameters
----------
device : torch.device, optional
The target device of the matrix if given, otherwise the current
The target device of the matrix if provided, otherwise the current
device will be used
dtype : torch.dtype, optional
The target data type of the matrix values if given, otherwise the
The target data type of the matrix values if provided, otherwise the
current data type will be used
Returns
-------
DiagMatrix
The result matrix
The converted matrix
Example
--------
......@@ -205,7 +198,7 @@ class DiagMatrix:
return diag(self.val.to(device=device, dtype=dtype), self.shape)
def cuda(self):
"""Move the matrix to GPU. If the matrix is already on GPU, the
"""Moves the matrix to GPU. If the matrix is already on GPU, the
original matrix will be returned. If multiple GPU devices exist,
'cuda:0' will be selected.
......@@ -226,7 +219,7 @@ class DiagMatrix:
return self.to(device="cuda")
def cpu(self):
"""Move the matrix to CPU. If the matrix is already on CPU, the
"""Moves the matrix to CPU. If the matrix is already on CPU, the
original matrix will be returned.
Returns
......@@ -246,7 +239,7 @@ class DiagMatrix:
return self.to(device="cpu")
def float(self):
"""Convert the matrix values to float data type. If the matrix already
"""Converts the matrix values to float data type. If the matrix already
uses float data type, the original matrix will be returned.
Returns
......@@ -266,7 +259,7 @@ class DiagMatrix:
return self.to(dtype=torch.float)
def double(self):
"""Convert the matrix values to double data type. If the matrix already
"""Converts the matrix values to double data type. If the matrix already
uses double data type, the original matrix will be returned.
Returns
......@@ -286,7 +279,7 @@ class DiagMatrix:
return self.to(dtype=torch.double)
def int(self):
"""Convert the matrix values to int data type. If the matrix already
"""Converts the matrix values to int data type. If the matrix already
uses int data type, the original matrix will be returned.
Returns
......@@ -306,7 +299,7 @@ class DiagMatrix:
return self.to(dtype=torch.int)
def long(self):
"""Convert the matrix values to long data type. If the matrix already
"""Converts the matrix values to long data type. If the matrix already
uses long data type, the original matrix will be returned.
Returns
......@@ -329,15 +322,15 @@ class DiagMatrix:
def diag(
val: torch.Tensor, shape: Optional[Tuple[int, int]] = None
) -> DiagMatrix:
"""Create a diagonal matrix based on the diagonal values
"""Creates a diagonal matrix based on the diagonal values.
Parameters
----------
val : torch.Tensor
Diagonal of the matrix. It can take shape (N) or (N, D).
Diagonal of the matrix, in 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)`.
If specified, :attr:`len(val)` must be equal to :attr:`min(shape)`,
otherwise, it will be inferred from :attr:`val`, i.e., ``(N, N)``
Returns
-------
......@@ -383,7 +376,7 @@ def identity(
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
) -> DiagMatrix:
"""Create a diagonal matrix with ones on the diagonal and zeros elsewhere
"""Creates a diagonal matrix with ones on the diagonal and zeros elsewhere.
Parameters
----------
......@@ -447,7 +440,8 @@ def identity(
def _diag_matrix_str(spmat: DiagMatrix) -> str:
"""Internal function for converting a diagonal matrix to string
representation."""
representation.
"""
values_str = str(spmat.val)
meta_str = f"size={spmat.shape}"
if spmat.val.dim() > 1:
......
python/dgl/sparse/sparse_matrix.py
View file @ fa74ea70
......@@ -11,128 +11,127 @@ class SparseMatrix:
def __init__(self, c_sparse_matrix: torch.ScriptObject):
self.c_sparse_matrix = c_sparse_matrix
def __repr__(self):
return _sparse_matrix_str(self)
@property
def val(self) -> torch.Tensor:
"""Get the values of the nonzero elements.
"""Returns the values of the non-zero elements.
Returns
-------
torch.Tensor
Values of the nonzero elements
Values of the non-zero elements
"""
return self.c_sparse_matrix.val()
@property
def shape(self) -> Tuple[int]:
"""Shape of the sparse matrix.
"""Returns the shape of the sparse matrix.
Returns
-------
Tuple[int]
The shape of the matrix
The shape of the sparse matrix
"""
return tuple(self.c_sparse_matrix.shape())
@property
def nnz(self) -> int:
"""The number of nonzero elements of the sparse matrix.
"""Returns the number of non-zero elements in the sparse matrix.
Returns
-------
int
The number of nonzero elements of the matrix
The number of non-zero elements of the matrix
"""
return self.c_sparse_matrix.nnz()
@property
def dtype(self) -> torch.dtype:
"""Data type of the values of the sparse matrix.
"""Returns the data type of the sparse matrix.
Returns
-------
torch.dtype
Data type of the values of the matrix
Data type of the sparse matrix
"""
# FIXME: find a proper way to pass dtype from C++ to Python
return self.c_sparse_matrix.val().dtype
@property
def device(self) -> torch.device:
"""Device of the sparse matrix.
"""Returns the device the sparse matrix is on.
Returns
-------
torch.device
Device of the matrix
The device the sparse matrix is on
"""
return self.c_sparse_matrix.device()
@property
def row(self) -> torch.Tensor:
"""Get the row indices of the nonzero elements.
"""Returns the row indices of the non-zero elements.
Returns
-------
tensor
Row indices of the nonzero elements
Row indices of the non-zero elements
"""
return self.coo()[0]
@property
def col(self) -> torch.Tensor:
"""Get the column indices of the nonzero elements.
"""Returns the column indices of the non-zero elements.
Returns
-------
tensor
Column indices of the nonzero elements
Column indices of the non-zero elements
"""
return self.coo()[1]
def __repr__(self):
return _sparse_matrix_str(self)
def coo(self) -> Tuple[torch.Tensor, ...]:
"""Get the coordinate (COO) representation of the sparse matrix.
"""Returns the coordinate (COO) representation of the sparse matrix.
Returns
-------
Tuple[torch.Tensor, torch.Tensor]
A tuple of tensors containing row and column coordinates.
A tuple of tensors containing row and column coordinates
"""
return self.c_sparse_matrix.coo()
def csr(self) -> Tuple[torch.Tensor, ...]:
r"""Get the compressed sparse row (CSR) representation of the sparse
r"""Returns the compressed sparse row (CSR) representation of the sparse
matrix.
Returns
-------
Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
A tuple of tensors containing row, column coordinates and value
indices.
indices
"""
return self.c_sparse_matrix.csr()
def csc(self) -> Tuple[torch.Tensor, ...]:
r"""Get the compressed sparse column (CSC) representation of the sparse
matrix.
r"""Returns the compressed sparse column (CSC) representation of the
sparse matrix.
Returns
-------
Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
A tuple of tensors containing row, column coordinates and value
indices.
indices
"""
return self.c_sparse_matrix.csc()
def to_dense(self) -> torch.Tensor:
"""Return a dense representation of the matrix.
"""Returns a copy in dense matrix format of the sparse matrix.
Returns
-------
torch.Tensor
Dense representation of the sparse matrix.
The copy in dense matrix format
"""
row, col = self.coo()
val = self.val
......@@ -175,7 +174,7 @@ class SparseMatrix:
return SparseMatrix(self.c_sparse_matrix.transpose())
def to(self, device=None, dtype=None):
"""Perform matrix dtype and/or device conversion. If the target device
"""Performs matrix dtype and/or device conversion. If the target device
and dtype are already in use, the original matrix will be returned.
Parameters
......@@ -184,13 +183,13 @@ class SparseMatrix:
The target device of the matrix if provided, otherwise the current
device will be used
dtype : torch.dtype, optional
The target data type of the matrix values, otherwise the current
data type will be used
The target data type of the matrix values if provided, otherwise the
current data type will be used
Returns
-------
SparseMatrix
The result matrix
The converted matrix
Example
--------
......@@ -224,7 +223,7 @@ class SparseMatrix:
return from_coo(row, col, val, self.shape)
def cuda(self):
"""Move the matrix to GPU. If the matrix is already on GPU, the
"""Moves the matrix to GPU. If the matrix is already on GPU, the
original matrix will be returned. If multiple GPU devices exist,
'cuda:0' will be selected.
......@@ -248,7 +247,7 @@ class SparseMatrix:
return self.to(device="cuda")
def cpu(self):
"""Move the matrix to CPU. If the matrix is already on CPU, the
"""Moves the matrix to CPU. If the matrix is already on CPU, the
original matrix will be returned.
Returns
......@@ -271,7 +270,7 @@ class SparseMatrix:
return self.to(device="cpu")
def float(self):
"""Convert the matrix values to float data type. If the matrix already
"""Converts the matrix values to float data type. If the matrix already
uses float data type, the original matrix will be returned.
Returns
......@@ -295,7 +294,7 @@ class SparseMatrix:
return self.to(dtype=torch.float)
def double(self):
"""Convert the matrix values to double data type. If the matrix already
"""Converts the matrix values to double data type. If the matrix already
uses double data type, the original matrix will be returned.
Returns
......@@ -318,7 +317,7 @@ class SparseMatrix:
return self.to(dtype=torch.double)
def int(self):
"""Convert the matrix values to int data type. If the matrix already
"""Converts the matrix values to int data type. If the matrix already
uses int data type, the original matrix will be returned.
Returns
......@@ -341,7 +340,7 @@ class SparseMatrix:
return self.to(dtype=torch.int)
def long(self):
"""Convert the matrix values to long data type. If the matrix already
"""Converts the matrix values to long data type. If the matrix already
uses long data type, the original matrix will be returned.
Returns
......@@ -364,14 +363,14 @@ class SparseMatrix:
return self.to(dtype=torch.long)
def coalesce(self):
"""Return a coalesced sparse matrix.
"""Returns 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
The coalescing process will accumulate the non-zero elements of the same
indices by summation.
The function does not support autograd.
......@@ -379,7 +378,7 @@ class SparseMatrix:
Returns
-------
SparseMatrix
The coalesced sparse matrix.
The coalesced sparse matrix
Examples
--------
......@@ -397,12 +396,7 @@ class SparseMatrix:
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.
"""Returns ``True`` if the sparse matrix contains duplicate indices.
Examples
--------
......@@ -424,14 +418,14 @@ def from_coo(
val: Optional[torch.Tensor] = None,
shape: Optional[Tuple[int, int]] = None,
) -> SparseMatrix:
"""Create a sparse matrix from row and column coordinates.
"""Creates 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.
......@@ -492,11 +486,11 @@ def from_csr(
val: Optional[torch.Tensor] = None,
shape: Optional[Tuple[int, int]] = None,
) -> SparseMatrix:
"""Create a sparse matrix from CSR indices.
"""Creates a sparse matrix from CSR indices.
For row i of the sparse matrix
- the column indices of the nonzero entries are stored in
- the column indices of the non-zero elements are stored in
``indices[indptr[i]: indptr[i+1]]``
- the corresponding values are stored in ``val[indptr[i]: indptr[i+1]]``
......@@ -504,9 +498,9 @@ def from_csr(
----------
indptr : tensor
Pointer to the column indices of shape (N + 1), where N is the number
of rows.
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.
......@@ -576,11 +570,11 @@ def from_csc(
val: Optional[torch.Tensor] = None,
shape: Optional[Tuple[int, int]] = None,
) -> SparseMatrix:
"""Create a sparse matrix from CSC indices.
"""Creates a sparse matrix from CSC indices.
For column i of the sparse matrix
- the row indices of the nonzero entries are stored in
- the row indices of the non-zero elements are stored in
``indices[indptr[i]: indptr[i+1]]``
- the corresponding values are stored in ``val[indptr[i]: indptr[i+1]]``
......@@ -588,9 +582,9 @@ def from_csc(
----------
indptr : tensor
Pointer to the row indices of shape N + 1, where N is the
number of columns.
number of columns
indices : tensor
The row indices of shape nnz.
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.
......@@ -655,17 +649,17 @@ def from_csc(
def val_like(mat: SparseMatrix, val: torch.Tensor) -> SparseMatrix:
"""Create a sparse matrix from an existing sparse matrix using new values.
"""Creates a sparse matrix from an existing sparse matrix using new values.
The new sparse matrix will have the same nonzero indices as the given
sparse matrix and use the given values as the new nonzero values.
The new sparse matrix will have the same non-zero indices as the given
sparse matrix and use the given values as the new non-zero values.
Parameters
----------
mat : SparseMatrix
An existing sparse matrix with nnz nonzero values
An existing sparse matrix with non-zero values
val : tensor
The new nonzero values, a tensor of shape (nnz) or (nnz, D)
The new values of the non-zero elements, a tensor of shape (nnz) or (nnz, D)
Returns
-------
......@@ -691,7 +685,8 @@ def val_like(mat: SparseMatrix, val: torch.Tensor) -> SparseMatrix:
def _sparse_matrix_str(spmat: SparseMatrix) -> str:
"""Internal function for converting a sparse matrix to string
representation."""
representation.
"""
indices_str = str(torch.stack(spmat.coo()))
values_str = str(spmat.val)
meta_str = f"size={spmat.shape}, nnz={spmat.nnz}"
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment