mirrors / dmlc--dgl

项目文件夹

dmlc--dgl
文件
better-conda
T
Israt Nisa cf03592743 Move mock version of dgl_sparse library to DGL main repo (#4524) 
* init

* Add api doc for sparse library

* support op btwn matrices with differnt sparsity

* Fixed docstring

* addresses comments

* lint check

* change keyword format to fmt

Co-authored-by: Israt Nisa <nisisrat@amazon.com>
2022-09-14 12:08:20 -04:00

270 行
8.2 KiB
Python
原始文件 永久链接 Blame 文件历史

"""dgl elementwise operators for sparse matrix module."""
import torch
from .sp_matrix import SparseMatrix
__all__ = ['add', 'sub', 'mul', 'div', 'rdiv', 'power', 'rpower']
def add(A, B):
"""Elementwise addition.
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix
Sparse matrix
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
Case 1: Add two matrices of same sparsity structure
>>> rowA = torch.tensor([1, 0, 2, 7, 1])
>>> colA = torch.tensor([0, 49, 2, 1, 7])
>>> valA = torch.tensor([10, 20, 30, 40, 50])
>>> A = SparseMatrix(rowA, colA, valA, shape=(10, 50))
>>> A + A
SparseMatrix(indices=tensor([[ 0, 1, 1, 2, 7],
[49, 0, 7, 2, 1]]),
values=tensor([ 40, 20, 100, 60, 80]),
shape=(10, 50), nnz=5)
>>> w = torch.arange(1, len(rowA)+1)
>>> A + A(w)
SparseMatrix(indices=tensor([[ 0, 1, 1, 2, 7],
[49, 0, 7, 2, 1]]),
values=tensor([21, 12, 53, 34, 45]),
shape=(10, 50), nnz=5)
Case 2: Add two matrices of different sparsity structure
>>> rowB = torch.tensor([1, 9, 2, 7, 1, 1, 0])
>>> colB = torch.tensor([0, 1, 2, 1, 7, 11, 15])
>>> valB = torch.tensor([1, 2, 3, 4, 5, 6])
>>> B = SparseMatrix(rowB, colB, valB, shape=(10, 50))
>>> A + B
SparseMatrix(indices=tensor([[ 0, 1, 1, 1, 2, 7, 9],
[49, 0, 7, 11, 2, 1, 1]]),
values=tensor([20, 11, 55, 6, 33, 44, 2]),
shape=(10, 50), nnz=7)
"""
if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
' B {} are expected to match'.format(A.shape, B.shape)
C = A.adj + B.adj
return SparseMatrix(C.indices()[0], C.indices()[1], C.values(), C.shape)
raise RuntimeError('Elementwise addition between {} and {} is not ' \
'supported.'.format(type(A), type(B)))
def sub(A, B):
"""Elementwise subtraction.
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix
Sparse matrix
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> rowA = torch.tensor([1, 0, 2, 7, 1])
>>> colA = torch.tensor([0, 49, 2, 1, 7])
>>> valA = torch.tensor([10, 20, 30, 40, 50])
>>> A = SparseMatrix(rowA, colA, valA, shape=(10, 50))
>>> rowB = torch.tensor([1, 9, 2, 7, 1, 1])
>>> colB = torch.tensor([0, 1, 2, 1, 7, 11])
>>> valB = torch.tensor([1, 2, 3, 4, 5, 6])
>>> B = SparseMatrix(rowB, colB, valB, shape=(10, 50))
>>> A - B
SparseMatrix(indices=tensor([[ 0, 1, 1, 1, 2, 7, 9],
[49, 0, 7, 11, 2, 1, 1]]),
values=tensor([20, 9, 45, -6, 27, 36, -2]),
shape=(10, 50), nnz=7
"""
if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
' B {} are expected to match.'.format(A.shape, B.shape)
C = A.adj - B.adj
return SparseMatrix(C.indices()[0], C.indices()[1], C.values(), C.shape)
raise RuntimeError('Elementwise subtraction between {} and {} is not ' \
'supported.'.format(type(A), type(B)))
def mul(A, B):
"""Elementwise multiplication.
Parameters
----------
A : SparseMatrix or scalar
Sparse matrix or scalar value
B : SparseMatrix or scalar
Sparse matrix or scalar value.
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
Case 1: Elementwise multiplication between two sparse matrices
>>> rowA = torch.tensor([1, 0, 2, 7, 1])
>>> colA = torch.tensor([0, 49, 2, 1, 7])
>>> valA = torch.tensor([10, 20, 30, 40, 50])
>>> A = SparseMatrix(rowA, colA, valA, shape=(10, 50))
>>> rowB = torch.tensor([1, 9, 2, 7, 1, 1])
>>> colB = torch.tensor([0, 1, 2, 1, 7, 11])
>>> valB = torch.tensor([1, 2, 3, 4, 5, 6])
>>> B = SparseMatrix(rowB, colB, valB, shape=(10, 50))
>>> A * B
SparseMatrix(indices=tensor([[1, 1, 2, 7],
[0, 7, 2, 1]]),
values=tensor([ 10, 250, 90, 160]),
shape=(10, 50), nnz=4)
Case 2: Elementwise multiplication between sparse matrix and scalar
>>> v_scalar = 2.5
>>> A * v_scalar
SparseMatrix(indices=tensor([[ 0, 1, 1, 2, 7],
[49, 0, 7, 2, 1]]),
values=tensor([ 50., 25., 125., 75., 100.]),
shape=(8, 50), nnz=5)
>>> v_scalar * A
SparseMatrix(indices=tensor([[ 0, 1, 1, 2, 7],
[49, 0, 7, 2, 1]]),
values=tensor([ 50., 25., 125., 75., 100.]),
shape=(8, 50), nnz=5)
"""
if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
assert A.shape == B.shape, 'The shape of sparse matrix A {} and' \
' B {} are expected to match.'.format(A.shape, B.shape)
A = A.adj if isinstance(A, SparseMatrix) else A
B = B.adj if isinstance(B, SparseMatrix) else B
C = A * B
return SparseMatrix(C.indices()[0], C.indices()[1], C.values(), C.shape)
def div(A, B):
"""Elementwise division.
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix or scalar
Sparse matrix or scalar value.
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
Case 1: Elementwise division between two matrices of same sparsity (matrices
with different sparsity is not supported)
>>> rowA = torch.tensor([1, 0, 2, 7, 1])
>>> colA = torch.tensor([0, 49, 2, 1, 7])
>>> valA = torch.tensor([10, 20, 30, 40, 50])
>>> A = SparseMatrix(rowA, colA, valA, shape=(10, 50))
>>> w = torch.arange(1, len(rowA)+1)
>>> A/A(w)
SparseMatrix(indices=tensor([[ 0, 1, 1, 2, 7],
[49, 0, 7, 2, 1]]),
values=tensor([20.0000, 5.0000, 16.6667, 7.5000, 8.0000]),
shape=(8, 50), nnz=5)
Case 2: Elementwise multiplication between sparse matrix and scalar
>>> A / v_scalar
SparseMatrix(indices=tensor([[ 0, 1, 1, 2, 7],
[49, 0, 7, 2, 1]]),
values=tensor([ 8., 4., 20., 12., 16.]),
shape=(8, 50), nnz=5)
"""
if isinstance(A, SparseMatrix) and isinstance(B, SparseMatrix):
# same sparsity structure
if torch.equal(A.indices("COO"), B.indices("COO")):
return SparseMatrix(A.row, A.col, A.val / B.val, A.shape)
raise ValueError('Division between matrices of different sparsity is not supported')
C = A.adj/B
return SparseMatrix(C.indices()[0], C.indices()[1], C.values(), C.shape)
def rdiv(A, B):
"""Elementwise division.
Parameters
----------
A : scalar
scalar value
B : SparseMatrix
Sparse matrix
"""
raise RuntimeError('Elementwise division between {} and {} is not ' \
'supported.'.format(type(A), type(B)))
def power(A, B):
"""Elementwise power operation.
Parameters
----------
A : SparseMatrix
Sparse matrix
B : scalar
scalar value.
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> rowA = torch.tensor([1, 0, 2, 7, 1])
>>> colA = torch.tensor([0, 49, 2, 1, 7])
>>> valA = torch.tensor([10, 20, 30, 40, 50])
>>> A = SparseMatrix(rowA, colA, valA, shape=(10, 50))
>>> pow(A, 2.5)
SparseMatrix(indices=tensor([[ 0, 1, 1, 2, 7],
[49, 0, 7, 2, 1]]),
values=tensor([ 1788.8544, 316.2278, 17677.6699, 4929.5029, 10119.2881]),
shape=(8, 50), nnz=5)
"""
if isinstance(B, SparseMatrix):
raise RuntimeError('Power operation between two sparse matrices is not supported')
return SparseMatrix(A.row, A.col, torch.pow(A.val, B))
def rpower(A, B):
"""Elementwise power operation.
Parameters
----------
A : scalar
scalar value.
B : SparseMatrix
Sparse matrix.
"""
raise RuntimeError('Power operation between {} and {} is not ' \
'supported.'.format(type(A), type(B)))
SparseMatrix.__add__ = add
SparseMatrix.__radd__ = add
SparseMatrix.__sub__ = sub
SparseMatrix.__rsub__ = sub
SparseMatrix.__mul__ = mul
SparseMatrix.__rmul__ = mul
SparseMatrix.__truediv__ = div
SparseMatrix.__rtruediv__ = rdiv
SparseMatrix.__pow__ = power
SparseMatrix.__rpow__ = rpower
在新工单中引用 查看 Git Blame 复制永久链接