wehub-resource-sync
62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
"""DGL unary operators for sparse matrix module."""
|
|
from .sparse_matrix import diag, SparseMatrix, val_like
|
|
|
|
|
|
def neg(A: SparseMatrix) -> SparseMatrix:
|
|
"""Returns a new sparse matrix with the negation of the original nonzero
|
|
values, equivalent to ``-A``.
|
|
|
|
Returns
|
|
-------
|
|
SparseMatrix
|
|
Negation of the sparse matrix
|
|
|
|
Examples
|
|
--------
|
|
|
|
>>> indices = torch.tensor([[1, 1, 3], [1, 2, 3]])
|
|
>>> val = torch.tensor([1., 1., 2.])
|
|
>>> A = dglsp.spmatrix(indices, val)
|
|
>>> A = -A
|
|
SparseMatrix(indices=tensor([[1, 1, 3],
|
|
[1, 2, 3]]),
|
|
values=tensor([-1., -1., -2.]),
|
|
shape=(4, 4), nnz=3)
|
|
"""
|
|
return val_like(A, -A.val)
|
|
|
|
|
|
def inv(A: SparseMatrix) -> SparseMatrix:
|
|
"""Returns the inverse of the sparse matrix.
|
|
|
|
This function only supports square diagonal matrices with scalar nonzero
|
|
values.
|
|
|
|
Returns
|
|
-------
|
|
SparseMatrix
|
|
Inverse of the sparse matrix
|
|
|
|
Examples
|
|
--------
|
|
|
|
>>> val = torch.arange(1, 4).float()
|
|
>>> D = dglsp.diag(val)
|
|
>>> D.inv()
|
|
SparseMatrix(indices=tensor([[0, 1, 2],
|
|
[0, 1, 2]]),
|
|
values=tensor([1., 2., 3.]),
|
|
shape=(3, 3), nnz=3)
|
|
"""
|
|
num_rows, num_cols = A.shape
|
|
assert A.is_diag(), "Non-diagonal sparse matrix does not support inversion."
|
|
assert num_rows == num_cols, f"Expect a square matrix, got shape {A.shape}"
|
|
assert len(A.val.shape) == 1, "inv only supports 1D nonzero val"
|
|
|
|
return diag(1.0 / A.val, A.shape)
|
|
|
|
|
|
SparseMatrix.neg = neg
|
|
SparseMatrix.__neg__ = neg
|
|
SparseMatrix.inv = inv
|