Source code for dgl.sparse.sparse_matrix

"""DGL sparse matrix module."""
# pylint: disable= invalid-name
from typing import Optional, Tuple

import torch


[docs]class SparseMatrix: r"""Class for sparse matrix.""" 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: """Returns the values of the non-zero elements. Returns ------- torch.Tensor Values of the non-zero elements """ return self.c_sparse_matrix.val() @property def shape(self) -> Tuple[int]: """Returns the shape of the sparse matrix. Returns ------- Tuple[int] The shape of the sparse matrix """ return tuple(self.c_sparse_matrix.shape()) @property def nnz(self) -> int: """Returns the number of non-zero elements in the sparse matrix. Returns ------- int The number of non-zero elements of the matrix """ return self.c_sparse_matrix.nnz() @property def dtype(self) -> torch.dtype: """Returns the data type of the sparse matrix. Returns ------- torch.dtype Data type of the sparse matrix """ return self.c_sparse_matrix.val().dtype @property def device(self) -> torch.device: """Returns the device the sparse matrix is on. Returns ------- torch.device The device the sparse matrix is on """ return self.c_sparse_matrix.device() @property def row(self) -> torch.Tensor: """Returns the row indices of the non-zero elements. Returns ------- torch.Tensor Row indices of the non-zero elements """ return self.coo()[0] @property def col(self) -> torch.Tensor: """Returns the column indices of the non-zero elements. Returns ------- torch.Tensor Column indices of the non-zero elements """ return self.coo()[1]
[docs] def coo(self) -> Tuple[torch.Tensor, torch.Tensor]: r"""Returns the coordinate list (COO) representation of the sparse matrix. See `COO in Wikipedia <https://en.wikipedia.org/wiki/ Sparse_matrix#Coordinate_list_(COO)>`_. Returns ------- torch.Tensor Row coordinate torch.Tensor Column coordinate Examples -------- >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]]) >>> A = dglsp.spmatrix(indices) >>> A.coo() (tensor([1, 2, 1]), tensor([2, 4, 3])) """ return self.c_sparse_matrix.coo()
[docs] 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()
[docs] def csr(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: r"""Returns the compressed sparse row (CSR) representation of the sparse matrix. See `CSR in Wikipedia <https://en.wikipedia.org/wiki/ Sparse_matrix#Compressed_sparse_row_(CSR, _CRS_or_Yale_format)>`_. This function also returns value indices as an index tensor, indicating the order of the values of non-zero elements in the CSR representation. A ``None`` value indices array indicates the order of the values stays the same as the values of the SparseMatrix. Returns ------- torch.Tensor Row indptr torch.Tensor Column indices torch.Tensor Value indices Examples -------- >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]]) >>> A = dglsp.spmatrix(indices) >>> A.csr() (tensor([0, 0, 2, 3]), tensor([2, 3, 4]), tensor([0, 2, 1])) """ return self.c_sparse_matrix.csr()
[docs] def csc(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: r"""Returns the compressed sparse column (CSC) representation of the sparse matrix. See `CSC in Wikipedia <https://en.wikipedia.org/wiki/ Sparse_matrix#Compressed_sparse_column_(CSC_or_CCS)>`_. This function also returns value indices as an index tensor, indicating the order of the values of non-zero elements in the CSC representation. A ``None`` value indices array indicates the order of the values stays the same as the values of the SparseMatrix. Returns ------- torch.Tensor Column indptr torch.Tensor Row indices torch.Tensor Value indices Examples -------- >>> indices = torch.tensor([[1, 2, 1], [2, 4, 3]]) >>> A = dglsp.spmatrix(indices) >>> A.csc() (tensor([0, 0, 0, 1, 2, 3]), tensor([1, 1, 2]), tensor([0, 2, 1])) """ return self.c_sparse_matrix.csc()
[docs] def to_dense(self) -> torch.Tensor: """Returns a copy in dense matrix format of the sparse matrix. Returns ------- torch.Tensor The copy in dense matrix format """ row, col = self.coo() val = self.val shape = self.shape + val.shape[1:] mat = torch.zeros(shape, device=self.device, dtype=self.dtype) mat[row, col] = val return mat
[docs] def t(self): """Alias of :meth:`transpose()`""" return self.transpose()
@property def T(self): # pylint: disable=C0103 """Alias of :meth:`transpose()`""" return self.transpose()
[docs] def transpose(self): """Returns the transpose of this sparse matrix. Returns ------- SparseMatrix The transpose of this sparse matrix. Examples -------- >>> indices = torch.tensor([[1, 1, 3], [2, 1, 3]]) >>> val = torch.tensor([1, 1, 2]) >>> A = dglsp.spmatrix(indices, val) >>> A = A.transpose() SparseMatrix(indices=tensor([[2, 1, 3], [1, 1, 3]]), values=tensor([1, 1, 2]), shape=(4, 4), nnz=3) """ return SparseMatrix(self.c_sparse_matrix.transpose())
[docs] def to(self, device=None, dtype=None): """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 provided, otherwise the current device will be used dtype : torch.dtype, optional The target data type of the matrix values if provided, otherwise the current data type will be used Returns ------- SparseMatrix The converted matrix Examples -------- >>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]]) >>> A = dglsp.spmatrix(indices, shape=(3, 4)) >>> A.to(device="cuda:0", dtype=torch.int32) SparseMatrix(indices=tensor([[1, 1, 2], [1, 2, 0]], device='cuda:0'), values=tensor([1, 1, 1], device='cuda:0', dtype=torch.int32), shape=(3, 4), nnz=3) """ if device is None: device = self.device if dtype is None: dtype = self.dtype if device == self.device and dtype == self.dtype: return self elif device == self.device: return val_like(self, self.val.to(dtype=dtype)) else: # TODO(#5119): Find a better moving strategy instead of always # convert to COO format. row, col = self.coo() row = row.to(device=device) col = col.to(device=device) val = self.val.to(device=device, dtype=dtype) return from_coo(row, col, val, self.shape)
[docs] def cuda(self): """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. Returns ------- SparseMatrix The matrix on GPU Examples -------- >>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]]) >>> A = dglsp.spmatrix(indices, shape=(3, 4)) >>> A.cuda() SparseMatrix(indices=tensor([[1, 1, 2], [1, 2, 0]], device='cuda:0'), values=tensor([1., 1., 1.], device='cuda:0'), shape=(3, 4), nnz=3) """ return self.to(device="cuda")
[docs] def cpu(self): """Moves the matrix to CPU. If the matrix is already on CPU, the original matrix will be returned. Returns ------- SparseMatrix The matrix on CPU Examples -------- >>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]]).to("cuda") >>> A = dglsp.spmatrix(indices, shape=(3, 4)) >>> A.cpu() SparseMatrix(indices=tensor([[1, 1, 2], [1, 2, 0]]), values=tensor([1., 1., 1.]), shape=(3, 4), nnz=3) """ return self.to(device="cpu")
[docs] def float(self): """Converts the matrix values to float32 data type. If the matrix already uses float data type, the original matrix will be returned. Returns ------- SparseMatrix The matrix with float values Examples -------- >>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]]) >>> val = torch.ones(len(row)).long() >>> A = dglsp.spmatrix(indices, val, shape=(3, 4)) >>> A.float() SparseMatrix(indices=tensor([[1, 1, 2], [1, 2, 0]]), values=tensor([1., 1., 1.]), shape=(3, 4), nnz=3) """ return self.to(dtype=torch.float)
[docs] def double(self): """Converts the matrix values to double data type. If the matrix already uses double data type, the original matrix will be returned. Returns ------- SparseMatrix The matrix with double values Examples -------- >>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]]) >>> A = dglsp.spmatrix(indices, shape=(3, 4)) >>> A.double() SparseMatrix(indices=tensor([[1, 1, 2], [1, 2, 0]]), values=tensor([1., 1., 1.], dtype=torch.float64), shape=(3, 4), nnz=3) """ return self.to(dtype=torch.double)
[docs] def int(self): """Converts the matrix values to int32 data type. If the matrix already uses int data type, the original matrix will be returned. Returns ------- DiagMatrix The matrix with int values Examples -------- >>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]]) >>> A = dglsp.spmatrix(indices, shape=(3, 4)) >>> A.int() SparseMatrix(indices=tensor([[1, 1, 2], [1, 2, 0]]), values=tensor([1, 1, 1], dtype=torch.int32), shape=(3, 4), nnz=3) """ return self.to(dtype=torch.int)
[docs] def long(self): """Converts the matrix values to long data type. If the matrix already uses long data type, the original matrix will be returned. Returns ------- DiagMatrix The matrix with long values Examples -------- >>> indices = torch.tensor([[1, 1, 2], [1, 2, 0]]) >>> A = dglsp.spmatrix(indices, shape=(3, 4)) >>> A.long() SparseMatrix(indices=tensor([[1, 1, 2], [1, 2, 0]]), values=tensor([1, 1, 1]), shape=(3, 4), nnz=3) """ return self.to(dtype=torch.long)
[docs] def coalesce(self): """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 elements of the same indices by summation. The function does not support autograd. Returns ------- SparseMatrix The coalesced sparse matrix Examples -------- >>> indices = torch.tensor([[1, 0, 0, 0, 1], [1, 1, 1, 2, 2]]) >>> val = torch.tensor([0, 1, 2, 3, 4]) >>> A = dglsp.spmatrix(indices, val) >>> A.coalesce() SparseMatrix(indices=tensor([[0, 0, 1, 1], [1, 2, 1, 2]]), values=tensor([3, 3, 0, 4]), shape=(2, 3), nnz=4) """ return SparseMatrix(self.c_sparse_matrix.coalesce())
[docs] def has_duplicate(self): """Returns ``True`` if the sparse matrix contains duplicate indices. Examples -------- >>> indices = torch.tensor([[1, 0, 0, 0, 1], [1, 1, 1, 2, 2]]) >>> val = torch.tensor([0, 1, 2, 3, 4]) >>> A = dglsp.spmatrix(indices, val) >>> A.has_duplicate() True >>> A.coalesce().has_duplicate() False """ return self.c_sparse_matrix.has_duplicate()
def is_diag(self): """Returns whether the sparse matrix is a diagonal matrix.""" return self.c_sparse_matrix.is_diag() def index_select(self, dim: int, index: torch.Tensor): """Returns a sub-matrix selected according to the given index. Parameters ---------- dim : int The dim to select from matrix, should be 0 or 1. `dim = 0` for rowwise selection and `dim = 1` for columnwise selection. index : torch.Tensor The selection index indicates which IDs from the `dim` should be chosen from the matrix. Note that duplicated ids are allowed. The function does not support autograd. Returns ------- SparseMatrix The sub-matrix which contains selected rows or columns. Examples -------- >>> indices = torch.tensor([0, 1, 1, 2, 3, 4], [0, 2, 4, 3, 5, 0]]) >>> val = torch.tensor([0, 1, 2, 3, 4, 5]) >>> A = dglsp.spmatrix(indices, val) Case 1: Select rows by IDs. >>> row_ids = torch.tensor([0, 1, 4]) >>> A.index_select(0, row_ids) SparseMatrix(indices=tensor([[0, 1, 1, 2], [0, 2, 4, 0]]), values=tensor([0, 1, 2, 5]), shape=(3, 6), nnz=4) Case 2: Select columns by IDs. >>> column_ids = torch.tensor([0, 4, 5]) >>> A.index_select(1, column_ids) SparseMatrix(indices=tensor([[0, 4, 1, 3], [0, 0, 1, 2]]), values=tensor([0, 5, 2, 4]), shape=(5, 3), nnz=4) """ if dim not in (0, 1): raise ValueError("The selection dimension should be 0 or 1.") if isinstance(index, torch.Tensor): return SparseMatrix(self.c_sparse_matrix.index_select(dim, index)) raise TypeError(f"{type(index).__name__} is unsupported input type.") def range_select(self, dim: int, index: slice): """Returns a sub-matrix selected according to the given range index. Parameters ---------- dim : int The dim to select from matrix, should be 0 or 1. `dim = 0` for rowwise selection and `dim = 1` for columnwise selection. index : slice The selection slice indicates ID index from the `dim` should be chosen from the matrix. The function does not support autograd. Returns ------- SparseMatrix The sub-matrix which contains selected rows or columns. Examples -------- >>> indices = torch.tensor([0, 1, 1, 2, 3, 4], [0, 2, 4, 3, 5, 0]]) >>> val = torch.tensor([0, 1, 2, 3, 4, 5]) >>> A = dglsp.spmatrix(indices, val) Case 1: Select rows with given slice object. >>> A.range_select(0, slice(1, 3)) SparseMatrix(indices=tensor([[0, 0, 1], [2, 4, 3]]), values=tensor([1, 2, 3]), shape=(2, 6), nnz=3) Case 2: Select columns with given slice object. >>> A.range_select(1, slice(3, 6)) SparseMatrix(indices=tensor([[2, 1, 3], [0, 1, 2]]), values=tensor([3, 2, 4]), shape=(5, 3), nnz=3) """ if dim not in (0, 1): raise ValueError("The selection dimension should be 0 or 1.") if isinstance(index, slice): if index.step not in (None, 1): raise NotImplementedError( "Slice with step other than 1 are not supported yet." ) start = 0 if index.start is None else index.start end = index.stop return SparseMatrix( self.c_sparse_matrix.range_select(dim, start, end) ) raise TypeError(f"{type(index).__name__} is unsupported input type.") def sample( self, dim: int, fanout: int, ids: Optional[torch.Tensor] = None, replace: Optional[bool] = False, bias: Optional[bool] = False, ): """Returns a sampled matrix on the given dimension and sample arguments. Parameters ---------- dim : int The dimension for sampling, should be 0 or 1. `dim = 0` for rowwise selection and `dim = 1` for columnwise selection. fanout : int The number of elements to randomly sample on each row or column. ids : torch.Tensor, optional An optional tensor containing row or column IDs from which to sample elements. NOTE: If `ids` is not provided (i.e., `ids = None`), the function will sample from all rows or columns. replace : bool, optional Indicates whether repeated sampling of the same element is allowed. When `replace = True`, repeated sampling is permitted; when `replace = False`, it is not allowed. NOTE: If `replace = False` and there are fewer elements than `fanout`, all non-zero elements will be sampled. bias : bool, optional A boolean flag indicating whether to enable biasing during sampling. When `bias = True`, the values of the sparse matrix will be used as bias weights. The function does not support autograd. Returns ------- SparseMatrix A submatrix with the same shape as the original matrix, containing the randomly sampled non-zero elements. Examples -------- >>> indices = torch.tensor([[0, 0, 1, 1, 2, 2, 2], [0, 2, 0, 1, 0, 1, 2]]) >>> val = torch.tensor([0, 1, 2, 3, 4, 5, 6]) >>> A = dglsp.spmatrix(indices, val) Case 1: Sample rows with the given number and disable repeated sampling. >>> row_ids = torch.tensor([0, 2]) >>> A.sample(0, 2, row_ids) SparseMatrix(indices=tensor([[0, 0, 1, 1], [0, 2, 0, 2]]), values=tensor([0, 1, 4, 6]), shape=(2, 3), nnz=4) Case 2: Sample cols with the given number and disable repeated sampling. >>> col_ids = torch.tensor([0, 2]) >>> A.sample(1, 2, col_ids) SparseMatrix(indices=tensor([[0, 1, 0, 2], [0, 0, 1, 1]]), values=tensor([0, 2, 1, 6]), shape=(3, 2), nnz=4) Case 3: Sample rows with the given number and enable repeated sampling. >>> row_ids = torch.tensor([0, 1]) >>> A.sample(0, 2, row_ids, True) SparseMatrix(indices=tensor([[0, 0, 1, 1], [0, 2, 0, 0]]), values=tensor([0, 1, 2, 2]), shape=(2, 3), nnz=3) Case 4: Sample cols with the given number and enable repeated sampling. >>> col_ids = torch.tensor([0, 1]) >>> A.sample(1, 2, col_ids, True) SparseMatrix(indices=tensor([[0, 1, 1, 1], [0, 0, 1, 1]]), values=tensor([0, 2, 3, 3]), shape=(3, 2), nnz=3) """ if ids is None: dim_size = self.shape[0] if dim == 0 else self.shape[1] ids = torch.range( 0, dim_size, dtype=torch.int64, device=self.device ) return SparseMatrix( self.c_sparse_matrix.sample(dim, fanout, ids, replace, bias) ) def compact( self, dim: int, leading_indices: Optional[torch.Tensor] = None, ): """Compact sparse matrix by removing rows or columns without non-zero elements in the sparse matrix and relabeling indices of the dimension. This function serves a dual purpose: it allows you to reorganize the indices within a specific dimension (rows or columns) of the sparse matrix and, if needed, place certain 'leading_indices' at the beginning of the relabeled dimension. In the absence of 'leading_indices' (when it's set to `None`), the order of relabeled indices remains the same as the original order, except that rows or columns without non-zero elements are removed. When 'leading_indices' are provided, they are positioned at the start of the relabeled dimension. To be precise, all rows selected by the specified indices will be remapped from 0 to length(indices) - 1. Rows that are not selected and contain any non-zero elements will be positioned after those remapped rows while maintaining their original order. This function mimics 'dgl.to_block', a method used to compress a sampled subgraph by eliminating redundant nodes. The 'leading_indices' parameter replicates the behavior of 'include_dst_in_src' in 'dgl.to_block', adding destination node information for message passing. Setting 'leading_indices' to column IDs when relabeling the row dimension, for example, achieves the same effect as including destination nodes in source nodes. Parameters ---------- dim : int The dimension to relabel. Should be 0 or 1. Use `dim = 0` for rowwise relabeling and `dim = 1` for columnwise relabeling. leading_indices : torch.Tensor, optional An optional tensor containing row or column ids that should be placed at the beginning of the relabeled dimension. Returns ------- Tuple[SparseMatrix, torch.Tensor] A tuple containing the relabeled sparse matrix and the index mapping of the relabeled dimension from the new index to the original index. Examples -------- >>> indices = torch.tensor([[0, 2], [1, 2]]) >>> A = dglsp.spmatrix(indices) >>> print(A.to_dense()) tensor([[0., 1., 0.], [0., 0., 0.], [0., 0., 1.]]) Case 1: Compact rows without indices. >>> B, original_rows = A.compact(dim=0, leading_indices=None) >>> print(B.to_dense()) tensor([[0., 1., 0.], [0., 0., 1.]]) >>> print(original_rows) torch.Tensor([0, 2]) Case 2: Compact rows with indices. >>> B, original_rows = A.compact(dim=0, leading_indices=[1, 2]) >>> print(B.to_dense()) tensor([[0., 0., 0.], [0., 0., 1.], [0., 1., 0.],]) >>> print(original_rows) torch.Tensor([1, 2, 0]) """ mat, idx = torch.ops.dgl_sparse.compact( self.c_sparse_matrix, dim, leading_indices ) return SparseMatrix(mat), idx
[docs]def spmatrix( indices: torch.Tensor, val: Optional[torch.Tensor] = None, shape: Optional[Tuple[int, int]] = None, ) -> SparseMatrix: r"""Creates a sparse matrix from Coordinate format indices. Parameters ---------- indices : tensor.Tensor The indices are the coordinates of the non-zero elements in the matrix, which should have shape of ``(2, N)`` where the first row is the row indices and the second row is the column indices of non-zero elements. val : tensor.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 ------- SparseMatrix Sparse matrix Examples -------- Case1: Sparse matrix with row and column indices without values. >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]]) >>> A = dglsp.spmatrix(indices) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([1., 1., 1.]), shape=(3, 5), nnz=3) >>> # Specify shape >>> A = dglsp.spmatrix(indices, shape=(5, 5)) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([1., 1., 1.]), shape=(5, 5), nnz=3) Case2: Sparse matrix with scalar values. >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]]) >>> val = torch.tensor([[1.], [2.], [3.]]) >>> A = dglsp.spmatrix(indices, val) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([[1.], [2.], [3.]]), shape=(3, 5), nnz=3, val_size=(1,)) Case3: Sparse matrix with vector values. >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]]) >>> val = torch.tensor([[1., 1.], [2., 2.], [3., 3.]]) >>> A = dglsp.spmatrix(indices, val) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([[1., 1.], [2., 2.], [3., 3.]]), shape=(3, 5), nnz=3, val_size=(2,)) """ 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))
[docs]def from_coo( row: torch.Tensor, col: torch.Tensor, val: Optional[torch.Tensor] = None, shape: Optional[Tuple[int, int]] = None, ) -> SparseMatrix: r"""Creates a sparse matrix from a coordinate list (COO), which stores a list of (row, column, value) tuples. See `COO in Wikipedia <https://en.wikipedia.org/wiki/Sparse_matrix#Coordinate_list_(COO)>`_. Parameters ---------- row : torch.Tensor The row indices of shape ``(nnz)`` col : torch.Tensor The column indices of shape ``(nnz)`` val : torch.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 ------- SparseMatrix Sparse matrix Examples -------- Case1: Sparse matrix with row and column indices without values. >>> dst = torch.tensor([1, 1, 2]) >>> src = torch.tensor([2, 4, 3]) >>> A = dglsp.from_coo(dst, src) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([1., 1., 1.]), shape=(3, 5), nnz=3) >>> # Specify shape >>> A = dglsp.from_coo(dst, src, shape=(5, 5)) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([1., 1., 1.]), shape=(5, 5), nnz=3) Case2: Sparse matrix with scalar values. >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]]) >>> val = torch.tensor([[1.], [2.], [3.]]) >>> A = dglsp.spmatrix(indices, val) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([[1.], [2.], [3.]]), shape=(3, 5), nnz=3, val_size=(1,)) Case3: Sparse matrix with vector values. >>> dst = torch.tensor([1, 1, 2]) >>> src = torch.tensor([2, 4, 3]) >>> val = torch.tensor([[1., 1.], [2., 2.], [3., 3.]]) >>> A = dglsp.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, val_size=(2,)) """ assert row.shape[0] == col.shape[0] return spmatrix(torch.stack([row, col]), val, shape)
[docs]def from_csr( indptr: torch.Tensor, indices: torch.Tensor, val: Optional[torch.Tensor] = None, shape: Optional[Tuple[int, int]] = None, ) -> SparseMatrix: r"""Creates a sparse matrix from compress sparse row (CSR) format. See `CSR in Wikipedia <https://en.wikipedia.org/wiki/ Sparse_matrix#Compressed_sparse_row_(CSR,_CRS_or_Yale_format)>`_. For row i of the sparse matrix - 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]]`` Parameters ---------- indptr : torch.Tensor Pointer to the column indices of shape ``(N + 1)``, where ``N`` is the number of rows indices : torch.Tensor The column indices of shape ``(nnz)`` val : torch.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 ------- SparseMatrix Sparse matrix Examples -------- Case1: Sparse matrix without values .. code:: [[0, 1, 0], [0, 0, 1], [1, 1, 1]] >>> indptr = torch.tensor([0, 1, 2, 5]) >>> indices = torch.tensor([1, 2, 0, 1, 2]) >>> A = dglsp.from_csr(indptr, indices) 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 = dglsp.from_csr(indptr, indices, shape=(3, 5)) SparseMatrix(indices=tensor([[0, 1, 2, 2, 2], [1, 2, 0, 1, 2]]), values=tensor([1., 1., 1., 1., 1.]), shape=(3, 5), nnz=5) Case2: Sparse matrix with scalar/vector values. Following example is with vector data. >>> indptr = torch.tensor([0, 1, 2, 5]) >>> indices = torch.tensor([1, 2, 0, 1, 2]) >>> val = torch.tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]) >>> A = dglsp.from_csr(indptr, indices, val) SparseMatrix(indices=tensor([[0, 1, 2, 2, 2], [1, 2, 0, 1, 2]]), values=tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]), shape=(3, 3), nnz=5, val_size=(2,)) """ if shape is None: shape = (indptr.shape[0] - 1, torch.max(indices) + 1) if val is None: val = torch.ones(indices.shape[0]).to(indptr.device) assert ( val.dim() <= 2 ), "The values of a SparseMatrix can only be scalars or vectors." return SparseMatrix( torch.ops.dgl_sparse.from_csr(indptr, indices, val, shape) )
[docs]def from_csc( indptr: torch.Tensor, indices: torch.Tensor, val: Optional[torch.Tensor] = None, shape: Optional[Tuple[int, int]] = None, ) -> SparseMatrix: r"""Creates a sparse matrix from compress sparse column (CSC) format. See `CSC in Wikipedia <https://en.wikipedia.org/wiki/ Sparse_matrix#Compressed_sparse_column_(CSC_or_CCS)>`_. For column i of the sparse matrix - 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]]`` Parameters ---------- indptr : torch.Tensor Pointer to the row indices of shape N + 1, where N is the number of columns indices : torch.Tensor The row indices of shape nnz val : torch.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 ------- SparseMatrix Sparse matrix Examples -------- Case1: Sparse matrix without values .. code:: [[0, 1, 0], [0, 0, 1], [1, 1, 1]] >>> indptr = torch.tensor([0, 1, 3, 5]) >>> indices = torch.tensor([2, 0, 2, 1, 2]) >>> A = dglsp.from_csc(indptr, indices) 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 = dglsp.from_csc(indptr, indices, shape=(5, 3)) 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) Case2: Sparse matrix with scalar/vector values. Following example is with vector data. >>> indptr = torch.tensor([0, 1, 3, 5]) >>> indices = torch.tensor([2, 0, 2, 1, 2]) >>> val = torch.tensor([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]) >>> A = dglsp.from_csc(indptr, indices, val) 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, val_size=(2,)) """ if shape is None: shape = (torch.max(indices) + 1, indptr.shape[0] - 1) if val is None: val = torch.ones(indices.shape[0]).to(indptr.device) assert ( val.dim() <= 2 ), "The values of a SparseMatrix can only be scalars or vectors." return SparseMatrix( torch.ops.dgl_sparse.from_csc(indptr, indices, val, shape) )
[docs]def val_like(mat: SparseMatrix, val: torch.Tensor) -> SparseMatrix: """Creates a sparse matrix from an existing sparse matrix using new 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 non-zero values val : torch.Tensor The new values of the non-zero elements, a tensor of shape ``(nnz)`` or ``(nnz, D)`` Returns ------- SparseMatrix New sparse matrix Examples -------- >>> indices = torch.tensor([[1, 1, 2], [2, 4, 3]]) >>> val = torch.ones(3) >>> A = dglsp.spmatrix(indices, val) >>> A = dglsp.val_like(A, torch.tensor([2, 2, 2])) SparseMatrix(indices=tensor([[1, 1, 2], [2, 4, 3]]), values=tensor([2, 2, 2]), shape=(3, 5), nnz=3) """ assert ( val.dim() <= 2 ), "The values of a SparseMatrix can only be scalars or vectors." return SparseMatrix(torch.ops.dgl_sparse.val_like(mat.c_sparse_matrix, val))
[docs]def diag( val: torch.Tensor, shape: Optional[Tuple[int, int]] = None ) -> SparseMatrix: """Creates a sparse matrix based on the diagonal values. Parameters ---------- val : torch.Tensor Diagonal of the matrix, in shape ``(N)`` or ``(N, D)`` shape : tuple[int, int], optional 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 ------- SparseMatrix Sparse matrix Examples -------- Case1: 5-by-5 diagonal matrix with scaler values on the diagonal >>> import torch >>> val = torch.ones(5) >>> dglsp.diag(val) 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) Case2: 5-by-10 diagonal matrix with scaler values on the diagonal >>> val = torch.ones(5) >>> dglsp.diag(val, shape=(5, 10)) SparseMatrix(indices=tensor([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]]), values=tensor([1., 1., 1., 1., 1.]), shape=(5, 10), nnz=5) Case3: 5-by-5 diagonal matrix with vector values on the diagonal >>> val = torch.randn(5, 3) >>> D = dglsp.diag(val) >>> D.shape (5, 5) >>> D.nnz 5 """ assert ( val.dim() <= 2 ), "The values of a DiagMatrix can only be scalars or vectors." len_val = len(val) if shape is not None: assert len_val == min(shape), ( f"Expect len(val) to be min(shape) for a diagonal matrix, got" f"{len_val} for len(val) and {shape} for shape." ) else: shape = (len_val, len_val) return SparseMatrix(torch.ops.dgl_sparse.from_diag(val, shape))
[docs]def identity( shape: Tuple[int, int], d: Optional[int] = None, dtype: Optional[torch.dtype] = None, device: Optional[torch.device] = None, ) -> SparseMatrix: r"""Creates a sparse 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 ------- SparseMatrix Sparse matrix Examples -------- Case1: 3-by-3 matrix with scaler diagonal values .. code:: [[1, 0, 0], [0, 1, 0], [0, 0, 1]] >>> dglsp.identity(shape=(3, 3)) SparseMatrix(indices=tensor([[0, 1, 2], [0, 1, 2]]), values=tensor([1., 1., 1.]), shape=(3, 3), nnz=3) Case2: 3-by-5 matrix with scaler diagonal values .. code:: [[1, 0, 0, 0, 0], [0, 1, 0, 0, 0], [0, 0, 1, 0, 0]] >>> dglsp.identity(shape=(3, 5)) SparseMatrix(indices=tensor([[0, 1, 2], [0, 1, 2]]), values=tensor([1., 1., 1.]), shape=(3, 5), nnz=3) Case3: 3-by-3 matrix with vector diagonal values >>> dglsp.identity(shape=(3, 3), d=2) SparseMatrix(indices=tensor([[0, 1, 2], [0, 1, 2]]), values=tensor([[1., 1.], [1., 1.], [1., 1.]]), shape=(3, 3), nnz=3, val_size=(2,)) """ 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)
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. """ indices_str = str(torch.stack(spmat.coo())) values_str = str(spmat.val) meta_str = f"shape={spmat.shape}, nnz={spmat.nnz}" if spmat.val.dim() > 1: val_size = tuple(spmat.val.shape[1:]) meta_str += f", val_size={val_size}" prefix = f"{type(spmat).__name__}(" def _add_indent(_str, indent): lines = _str.split("\n") lines = [lines[0]] + [" " * indent + line for line in lines[1:]] return "\n".join(lines) final_str = ( "indices=" + _add_indent(indices_str, len("indices=")) + ",\n" + "values=" + _add_indent(values_str, len("values=")) + ",\n" + meta_str + ")" ) final_str = prefix + _add_indent(final_str, len(prefix)) return final_str