wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/sparse/__init__.py
+43
View File
@@ -0,0 +1,43 @@
"""dgl sparse class."""
import os
import sys
import torch
from .._ffi import libinfo
from .broadcast import *
from .elementwise_op import *
from .elementwise_op_sp import *
from .matmul import *
from .reduction import * # pylint: disable=W0622
from .sddmm import *
from .softmax import *
from .sparse_matrix import *
from .unary_op import *
def load_dgl_sparse():
"""Load DGL C++ sparse library"""
version = torch.__version__.split("+", maxsplit=1)[0]
if sys.platform.startswith("linux"):
basename = f"libdgl_sparse_pytorch_{version}.so"
elif sys.platform.startswith("darwin"):
basename = f"libdgl_sparse_pytorch_{version}.dylib"
elif sys.platform.startswith("win"):
basename = f"dgl_sparse_pytorch_{version}.dll"
else:
raise NotImplementedError("Unsupported system: %s" % sys.platform)
dirname = os.path.dirname(libinfo.find_lib_path()[0])
path = os.path.join(dirname, "dgl_sparse", basename)
if not os.path.exists(path):
raise FileNotFoundError(f"Cannot find DGL C++ sparse library at {path}")
try:
torch.classes.load_library(path)
except Exception: # pylint: disable=W0703
raise ImportError("Cannot load DGL C++ sparse library")
load_dgl_sparse()
python/dgl/sparse/broadcast.py
+133
View File
@@ -0,0 +1,133 @@
"""DGL broadcast operator module."""
import operator
import torch
from .sparse_matrix import SparseMatrix, val_like
def sp_broadcast_v(A: SparseMatrix, v: torch.Tensor, op: str) -> SparseMatrix:
"""Broadcast operator for sparse matrix and vector.
:attr:`v` is broadcasted to the shape of :attr:`A` and then the operator is
applied on the non-zero values of :attr:`A`.
There are two cases regarding the shape of v:
1. :attr:`v` is a vector of shape ``(1, A.shape[1])`` or ``(A.shape[1])``.
In this case, :attr:`v` is broadcasted on the row dimension of :attr:`A`.
2. :attr:`v` is a vector of shape ``(A.shape[0], 1)``. In this case,
:attr:`v` is broadcasted on the column dimension of :attr:`A`.
If ``A.val`` takes shape ``(nnz, D)``, then :attr:`v` will be broadcasted on
the ``D`` dimension.
Parameters
----------
A: SparseMatrix
Sparse matrix
v: torch.Tensor
Vector
op: str
Operator in ["add", "sub", "mul", "truediv"]
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> v = torch.tensor([1, 2, 3, 4])
>>> dglsp.sp_broadcast_v(A, v, "add")
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([11, 24, 33]),
shape=(3, 4), nnz=3)
>>> v = torch.tensor([1, 2, 3]).view(-1, 1)
>>> dglsp.sp_broadcast_v(A, v, "add")
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([12, 21, 33]),
shape=(3, 4), nnz=3)
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([[10, 20], [30, 40], [50, 60]])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> v = torch.tensor([1, 2, 3]).view(-1, 1)
>>> dglsp.sp_broadcast_v(A, v, "sub")
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([[ 8, 18],
[29, 39],
[47, 57]]),
shape=(3, 4), nnz=3, val_size=(2,))
"""
op = getattr(operator, op)
if v.dim() == 1:
v = v.view(1, -1)
shape_error_message = (
f"Dimension mismatch for broadcasting. Got A.shape = {A.shape} and"
f"v.shape = {v.shape}."
)
assert v.dim() <= 2 and (1 in v.shape), shape_error_message
broadcast_dim = None
# v can be broadcasted to A if exactly one dimension of v is 1 and the other
# is the same as A.
for d, (dim1, dim2) in enumerate(zip(A.shape, v.shape)):
assert dim2 in (1, dim1), shape_error_message
if dim1 != dim2:
assert broadcast_dim is None, shape_error_message
broadcast_dim = d
# A and v has the same shape of (1, *) or (*, 1).
if broadcast_dim is None:
broadcast_dim = 0 if A.shape[0] == 1 else 1
if broadcast_dim == 0:
v = v.view(-1)[A.col]
else:
v = v.view(-1)[A.row]
if A.val.dim() > 1:
v = v.view(-1, 1)
ret_val = op(A.val, v)
return val_like(A, ret_val)
def sp_add_v(A: SparseMatrix, v: torch.Tensor) -> SparseMatrix:
"""Broadcast addition for sparse matrix and vector.
See the definition of :func:`sp_broadcast_v` for details.
"""
return sp_broadcast_v(A, v, "add")
def sp_sub_v(A: SparseMatrix, v: torch.Tensor) -> SparseMatrix:
"""Broadcast substraction for sparse matrix and vector.
See the definition of :func:`sp_broadcast_v` for details.
"""
return sp_broadcast_v(A, v, "sub")
def sp_mul_v(A: SparseMatrix, v: torch.Tensor) -> SparseMatrix:
"""Broadcast multiply for sparse matrix and vector.
See the definition of :func:`sp_broadcast_v` for details.
"""
return sp_broadcast_v(A, v, "mul")
def sp_div_v(A: SparseMatrix, v: torch.Tensor) -> SparseMatrix:
"""Broadcast division for sparse matrix and vector.
See the definition of :func:`sp_broadcast_v` for details.
"""
return sp_broadcast_v(A, v, "truediv")
python/dgl/sparse/elementwise_op.py
+201
View File
@@ -0,0 +1,201 @@
# pylint: disable=anomalous-backslash-in-string
"""DGL elementwise operator module."""
from typing import Union
from .sparse_matrix import SparseMatrix
from .utils import Scalar
__all__ = ["add", "sub", "mul", "div", "power"]
def add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
r"""Elementwise addition for ``SparseMatrix``, equivalent to ``A + B``.
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix
Sparse matrix
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 1, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val)
>>> B = dglsp.diag(torch.arange(1, 4))
>>> dglsp.add(A, B)
SparseMatrix(indices=tensor([[0, 0, 1, 1, 2],
[0, 1, 0, 1, 2]]),
values=tensor([1, 20, 10, 2, 33]),
shape=(3, 3), nnz=5)
"""
return A + B
def sub(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
r"""Elementwise subtraction for ``SparseMatrix``, equivalent to ``A - B``.
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix
Sparse matrix
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 1, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val)
>>> B = dglsp.diag(torch.arange(1, 4))
>>> dglsp.sub(A, B)
SparseMatrix(indices=tensor([[0, 0, 1, 1, 2],
[0, 1, 0, 1, 2]]),
values=tensor([-1, 20, 10, -2, 27]),
shape=(3, 3), nnz=5)
"""
return A - B
def mul(
A: Union[SparseMatrix, Scalar], B: Union[SparseMatrix, Scalar]
) -> SparseMatrix:
r"""Elementwise multiplication for ``SparseMatrix``, equivalent to
``A * B``.
If both :attr:`A` and :attr:`B` are sparse matrices, both of them should be
diagonal matrices.
Parameters
----------
A : SparseMatrix or Scalar
Sparse matrix or scalar value
B : SparseMatrix or Scalar
Sparse matrix or scalar value
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val)
>>> dglsp.mul(A, 2)
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([20, 40, 60]),
shape=(3, 4), nnz=3)
>>> D = dglsp.diag(torch.arange(1, 4))
>>> dglsp.mul(D, 2)
SparseMatrix(indices=tensor([[0, 1, 2],
[0, 1, 2]]),
values=tensor([2, 4, 6]),
shape=(3, 3), nnz=3)
>>> D = dglsp.diag(torch.arange(1, 4))
>>> dglsp.mul(D, D)
SparseMatrix(indices=tensor([[0, 1, 2],
[0, 1, 2]]),
values=tensor([1, 4, 9]),
shape=(3, 3), nnz=3)
"""
return A * B
def div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
r"""Elementwise division for ``SparseMatrix``, equivalent to ``A / B``.
If both :attr:`A` and :attr:`B` are sparse matrices, both of them should be
diagonal matrices.
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix or Scalar
Sparse matrix or scalar value
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> A = dglsp.diag(torch.arange(1, 4))
>>> B = dglsp.diag(torch.arange(10, 13))
>>> dglsp.div(A, B)
SparseMatrix(indices=tensor([[0, 1, 2],
[0, 1, 2]]),
values=tensor([0.1000, 0.1818, 0.2500]),
shape=(3, 3), nnz=3)
>>> A = dglsp.diag(torch.arange(1, 4))
>>> dglsp.div(A, 2)
SparseMatrix(indices=tensor([[0, 1, 2],
[0, 1, 2]]),
values=tensor([0.5000, 1.0000, 1.5000]),
shape=(3, 3), nnz=3)
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([1, 2, 3])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> dglsp.div(A, 2)
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([0.5000, 1.0000, 1.5000]),
shape=(3, 4), nnz=3)
"""
return A / B
def power(A: SparseMatrix, scalar: Scalar) -> SparseMatrix:
r"""Elementwise exponentiation ``SparseMatrix``, equivalent to
``A ** scalar``.
Parameters
----------
A : SparseMatrix
Sparse matrix
scalar : Scalar
Exponent
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val)
>>> dglsp.power(A, 2)
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([100, 400, 900]),
shape=(3, 4), nnz=3)
>>> D = dglsp.diag(torch.arange(1, 4))
>>> dglsp.power(D, 2)
SparseMatrix(indices=tensor([[0, 1, 2],
[0, 1, 2]]),
values=tensor([1, 4, 9]),
shape=(3, 3), nnz=3)
"""
return A**scalar
python/dgl/sparse/elementwise_op_sp.py
+220
View File
@@ -0,0 +1,220 @@
"""DGL elementwise operators for sparse matrix module."""
from typing import Union
import torch
from .sparse_matrix import SparseMatrix, val_like
from .utils import is_scalar, Scalar
def spsp_add(A, B):
"""Invoke C++ sparse library for addition"""
return SparseMatrix(
torch.ops.dgl_sparse.spsp_add(A.c_sparse_matrix, B.c_sparse_matrix)
)
def spsp_mul(A, B):
"""Invoke C++ sparse library for multiplication"""
return SparseMatrix(
torch.ops.dgl_sparse.spsp_mul(A.c_sparse_matrix, B.c_sparse_matrix)
)
def spsp_div(A, B):
"""Invoke C++ sparse library for division"""
return SparseMatrix(
torch.ops.dgl_sparse.spsp_div(A.c_sparse_matrix, B.c_sparse_matrix)
)
def sp_add(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
"""Elementwise addition
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix
Sparse matrix
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> A + A
SparseMatrix(indices=tensor([[0, 1, 2],
[3, 0, 2]]),
values=tensor([40, 20, 60]),
shape=(3, 4), nnz=3)
"""
# Python falls back to B.__radd__ then TypeError when NotImplemented is
# returned.
return spsp_add(A, B) if isinstance(B, SparseMatrix) else NotImplemented
def sp_sub(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
"""Elementwise subtraction
Parameters
----------
A : SparseMatrix
Sparse matrix
B : SparseMatrix
Sparse matrix
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> val2 = torch.tensor([5, 10, 15])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> B = dglsp.spmatrix(indices, val2, shape=(3, 4))
>>> A - B
SparseMatrix(indices=tensor([[0, 1, 2],
[3, 0, 2]]),
values=tensor([10, 5, 15]),
shape=(3, 4), nnz=3)
"""
# Python falls back to B.__rsub__ then TypeError when NotImplemented is
# returned.
return spsp_add(A, -B) if isinstance(B, SparseMatrix) else NotImplemented
def sp_mul(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
"""Elementwise multiplication
Note that if both :attr:`A` and :attr:`B` are sparse matrices, both of them
need to be diagonal or on CPU.
Parameters
----------
A : SparseMatrix
First operand
B : SparseMatrix or Scalar
Second operand
Returns
-------
SparseMatrix
Result of A * B
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([1, 2, 3])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> A * 2
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([2, 4, 6]),
shape=(3, 4), nnz=3)
>>> 2 * A
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([2, 4, 6]),
shape=(3, 4), nnz=3)
>>> indices2 = torch.tensor([[2, 0, 1], [0, 3, 2]])
>>> val2 = torch.tensor([3, 2, 1])
>>> B = dglsp.spmatrix(indices2, val2, shape=(3, 4))
>>> A * B
SparseMatrix(indices=tensor([[0],
[3]]),
values=tensor([4]),
shape=(3, 4), nnz=1)
"""
if is_scalar(B):
return val_like(A, A.val * B)
return spsp_mul(A, B)
def sp_div(A: SparseMatrix, B: Union[SparseMatrix, Scalar]) -> SparseMatrix:
"""Elementwise division
If :attr:`B` is a sparse matrix, both :attr:`A` and :attr:`B` must have the
same sparsity. And the returned matrix has the same order of non-zero
entries as :attr:`A`.
Parameters
----------
A : SparseMatrix
First operand
B : SparseMatrix or Scalar
Second operand
Returns
-------
SparseMatrix
Result of A / B
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([1, 2, 3])
>>> A = dglsp.spmatrix(indices, val, shape=(3, 4))
>>> A / 2
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([0.5000, 1.0000, 1.5000]),
shape=(3, 4), nnz=3)
"""
if is_scalar(B):
return val_like(A, A.val / B)
return spsp_div(A, B)
def sp_power(A: SparseMatrix, scalar: Scalar) -> SparseMatrix:
"""Take the power of each nonzero element and return a sparse matrix with
the result.
Parameters
----------
A : SparseMatrix
Sparse matrix
scalar : float or int
Exponent
Returns
-------
SparseMatrix
Sparse matrix
Examples
--------
>>> indices = torch.tensor([[1, 0, 2], [0, 3, 2]])
>>> val = torch.tensor([10, 20, 30])
>>> A = dglsp.spmatrix(indices, val)
>>> A ** 2
SparseMatrix(indices=tensor([[1, 0, 2],
[0, 3, 2]]),
values=tensor([100, 400, 900]),
shape=(3, 4), nnz=3)
"""
# Python falls back to scalar.__rpow__ then TypeError when NotImplemented
# is returned.
return val_like(A, A.val**scalar) if is_scalar(scalar) else NotImplemented
SparseMatrix.__add__ = sp_add
SparseMatrix.__sub__ = sp_sub
SparseMatrix.__mul__ = sp_mul
SparseMatrix.__rmul__ = sp_mul
SparseMatrix.__truediv__ = sp_div
SparseMatrix.__pow__ = sp_power
python/dgl/sparse/matmul.py
+227
View File
@@ -0,0 +1,227 @@
"""Matmul ops for SparseMatrix"""
# pylint: disable=invalid-name
from typing import Union
import torch
from .sparse_matrix import SparseMatrix
__all__ = ["spmm", "bspmm", "spspmm", "matmul"]
def spmm(A: SparseMatrix, X: torch.Tensor) -> torch.Tensor:
"""Multiplies a sparse matrix by a dense matrix, equivalent to ``A @ X``.
Parameters
----------
A : SparseMatrix
Sparse matrix of shape ``(L, M)`` with scalar values
X : torch.Tensor
Dense matrix of shape ``(M, N)`` or ``(M)``
Returns
-------
torch.Tensor
The dense matrix of shape ``(L, N)`` or ``(L)``
Examples
--------
>>> indices = torch.tensor([[0, 1, 1], [1, 0, 1]])
>>> val = torch.randn(indices.shape[1])
>>> A = dglsp.spmatrix(indices, val)
>>> X = torch.randn(2, 3)
>>> result = dglsp.spmm(A, X)
>>> type(result)
<class 'torch.Tensor'>
>>> result.shape
torch.Size([2, 3])
"""
assert isinstance(
A, SparseMatrix
), f"Expect arg1 to be a SparseMatrix object, got {type(A)}."
assert isinstance(
X, torch.Tensor
), f"Expect arg2 to be a torch.Tensor, got {type(X)}."
return torch.ops.dgl_sparse.spmm(A.c_sparse_matrix, X)
def bspmm(A: SparseMatrix, X: torch.Tensor) -> torch.Tensor:
"""Multiplies a sparse matrix by a dense matrix by batches, equivalent to
``A @ X``.
Parameters
----------
A : SparseMatrix
Sparse matrix of shape ``(L, M)`` with vector values of length ``K``
X : torch.Tensor
Dense matrix of shape ``(M, N, K)``
Returns
-------
torch.Tensor
Dense matrix of shape ``(L, N, K)``
Examples
--------
>>> indices = torch.tensor([[0, 1, 1], [1, 0, 2]])
>>> val = torch.randn(len(row), 2)
>>> A = dglsp.spmatrix(indices, val, shape=(3, 3))
>>> X = torch.randn(3, 3, 2)
>>> result = dglsp.bspmm(A, X)
>>> type(result)
<class 'torch.Tensor'>
>>> result.shape
torch.Size([3, 3, 2])
"""
assert isinstance(
A, SparseMatrix
), f"Expect arg1 to be a SparseMatrix object, got {type(A)}."
assert isinstance(
X, torch.Tensor
), f"Expect arg2 to be a torch.Tensor, got {type(X)}."
return spmm(A, X)
def spspmm(A: SparseMatrix, B: SparseMatrix) -> SparseMatrix:
"""Multiplies a sparse matrix by a sparse matrix, equivalent to ``A @ B``.
The non-zero values of the two sparse matrices must be 1D.
Parameters
----------
A : SparseMatrix
Sparse matrix of shape ``(L, M)``
B : SparseMatrix
Sparse matrix of shape ``(M, N)``
Returns
-------
SparseMatrix
Sparse matrix of shape ``(L, N)``.
Examples
--------
>>> indices1 = torch.tensor([[0, 1, 1], [1, 0, 1]])
>>> val1 = torch.ones(len(row1))
>>> A = dglsp.spmatrix(indices1, val1)
>>> indices2 = torch.tensor([[0, 1, 1], [0, 2, 1]])
>>> val2 = torch.ones(len(row2))
>>> B = dglsp.spmatrix(indices2, val2)
>>> dglsp.spspmm(A, B)
SparseMatrix(indices=tensor([[0, 0, 1, 1, 1],
[1, 2, 0, 1, 2]]),
values=tensor([1., 1., 1., 1., 1.]),
shape=(2, 3), nnz=5)
"""
assert isinstance(
A, SparseMatrix
), f"Expect A1 to be a SparseMatrix object, got {type(A)}."
assert isinstance(
B, SparseMatrix
), f"Expect A2 to be a SparseMatrix object, got {type(B)}."
return SparseMatrix(
torch.ops.dgl_sparse.spspmm(A.c_sparse_matrix, B.c_sparse_matrix)
)
def matmul(
A: Union[torch.Tensor, SparseMatrix], B: Union[torch.Tensor, SparseMatrix]
) -> Union[torch.Tensor, SparseMatrix]:
"""Multiplies two dense/sparse matrices, equivalent to ``A @ B``.
This function does not support the case where :attr:`A` is a \
``torch.Tensor`` and :attr:`B` is a ``SparseMatrix``.
* If both matrices are torch.Tensor, it calls \
:func:`torch.matmul()`. The result is a dense matrix.
* If both matrices are sparse, it calls :func:`dgl.sparse.spspmm`. The \
result is a sparse matrix.
* If :attr:`A` is sparse while :attr:`B` is dense, it calls \
:func:`dgl.sparse.spmm`. The result is a dense matrix.
* The operator supports batched sparse-dense matrix multiplication. In \
this case, the sparse matrix :attr:`A` should have shape ``(L, M)``, \
where the non-zero values have a batch dimension ``K``. The dense \
matrix :attr:`B` should have shape ``(M, N, K)``. The output \
is a dense matrix of shape ``(L, N, K)``.
* Sparse-sparse matrix multiplication does not support batched computation.
Parameters
----------
A : torch.Tensor or SparseMatrix
The first matrix.
B : torch.Tensor or SparseMatrix
The second matrix.
Returns
-------
torch.Tensor or SparseMatrix
The result matrix
Examples
--------
Multiplies a diagonal matrix with a dense matrix.
>>> val = torch.randn(3)
>>> A = dglsp.diag(val)
>>> B = torch.randn(3, 2)
>>> result = dglsp.matmul(A, B)
>>> type(result)
<class 'torch.Tensor'>
>>> result.shape
torch.Size([3, 2])
Multiplies a sparse matrix with a dense matrix.
>>> indices = torch.tensor([[0, 1, 1], [1, 0, 1]])
>>> val = torch.randn(indices.shape[1])
>>> A = dglsp.spmatrix(indices, val)
>>> X = torch.randn(2, 3)
>>> result = dglsp.matmul(A, X)
>>> type(result)
<class 'torch.Tensor'>
>>> result.shape
torch.Size([2, 3])
Multiplies a sparse matrix with a sparse matrix.
>>> indices1 = torch.tensor([[0, 1, 1], [1, 0, 1]])
>>> val1 = torch.ones(indices1.shape[1])
>>> A = dglsp.spmatrix(indices1, val1)
>>> indices2 = torch.tensor([[0, 1, 1], [0, 2, 1]])
>>> val2 = torch.ones(indices2.shape[1])
>>> B = dglsp.spmatrix(indices2, val2)
>>> result = dglsp.matmul(A, B)
>>> type(result)
<class 'dgl.sparse.sparse_matrix.SparseMatrix'>
>>> result.shape
(2, 3)
"""
assert isinstance(
A, (torch.Tensor, SparseMatrix)
), f"Expect arg1 to be a torch.Tensor or SparseMatrix, got {type(A)}."
assert isinstance(B, (torch.Tensor, SparseMatrix)), (
f"Expect arg2 to be a torch Tensor or SparseMatrix"
f"object, got {type(B)}."
)
if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
return torch.matmul(A, B)
assert not isinstance(A, torch.Tensor), (
f"Expect arg2 to be a torch Tensor if arg 1 is torch Tensor, "
f"got {type(B)}."
)
if isinstance(B, torch.Tensor):
return spmm(A, B)
return spspmm(A, B)
SparseMatrix.__matmul__ = matmul
python/dgl/sparse/reduction.py
+388
View File
@@ -0,0 +1,388 @@
"""DGL sparse matrix reduce operators"""
# pylint: disable=W0622
from typing import Optional
import torch
from .sparse_matrix import SparseMatrix
def reduce(input: SparseMatrix, dim: Optional[int] = None, rtype: str = "sum"):
"""Computes the reduction of non-zero values of the :attr:`input` sparse
matrix along the given dimension :attr:`dim`.
The reduction does not count zero elements. If the row or column to be
reduced does not have any non-zero elements, the result will be 0.
Parameters
----------
input : SparseMatrix
The input sparse matrix
dim : int, optional
The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on both rows and columns simultaneously)
If :attr:`dim` is None, it reduces both the rows and the columns
in the sparse matrix, producing a tensor of shape
``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
or column (``dim=1``) dimension, producing a tensor of shape
``(input.shape[1],) + input.val.shape[1:]`` or
``(input.shape[0],) + input.val.shape[1:]``.
rtype: str, optional
Reduction type, one of ``['sum', 'smin', 'smax', 'smean', 'sprod']``,
representing taking the sum, minimum, maximum, mean, and product of the
non-zero elements
Returns
----------
torch.Tensor
Reduced tensor
Examples
----------
Case1: scalar-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.reduce(A, rtype='sum')
tensor(4)
>>> dglsp.reduce(A, 0, 'sum')
tensor([2, 0, 2])
>>> dglsp.reduce(A, 1, 'sum')
tensor([1, 3, 0, 0])
>>> dglsp.reduce(A, 0, 'smax')
tensor([1, 0, 2])
>>> dglsp.reduce(A, 1, 'smin')
tensor([1, 1, 0, 0])
Case2: vector-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([[1., 2.], [2., 1.], [2., 2.]])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.reduce(A, rtype='sum')
tensor([5., 5.])
>>> dglsp.reduce(A, 0, 'sum')
tensor([[3., 3.],
[0., 0.],
[2., 2.]])
>>> dglsp.reduce(A, 1, 'smin')
tensor([[1., 2.],
[2., 1.],
[0., 0.],
[0., 0.]])
>>> dglsp.reduce(A, 0, 'smean')
tensor([[1.5000, 1.5000],
[0.0000, 0.0000],
[2.0000, 2.0000]])
"""
return torch.ops.dgl_sparse.reduce(input.c_sparse_matrix, rtype, dim)
def sum(input: SparseMatrix, dim: Optional[int] = None):
"""Computes the sum of non-zero values of the :attr:`input` sparse matrix
along the given dimension :attr:`dim`.
Parameters
----------
input : SparseMatrix
The input sparse matrix
dim : int, optional
The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on both rows and columns simultaneously)
If :attr:`dim` is None, it reduces both the rows and the columns
in the sparse matrix, producing a tensor of shape
``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
or column (``dim=1``) dimension, producing a tensor of shape
``(input.shape[1],) + input.val.shape[1:]`` or
``(input.shape[0],) + input.val.shape[1:]``.
Returns
----------
torch.Tensor
Reduced tensor
Examples
----------
Case1: scalar-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.sum(A)
tensor(4)
>>> dglsp.sum(A, 0)
tensor([2, 0, 2])
>>> dglsp.sum(A, 1)
tensor([1, 3, 0, 0])
Case2: vector-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.sum(A)
tensor([5, 5])
>>> dglsp.sum(A, 0)
tensor([[3, 3],
[0, 0],
[2, 2]])
"""
return torch.ops.dgl_sparse.sum(input.c_sparse_matrix, dim)
def smax(input: SparseMatrix, dim: Optional[int] = None):
"""Computes the maximum of non-zero values of the :attr:`input` sparse
matrix along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be
reduced does not have any non-zero value, the result will be 0.
Parameters
----------
input : SparseMatrix
The input sparse matrix
dim : int, optional
The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on both rows and columns simultaneously)
If :attr:`dim` is None, it reduces both the rows and the columns
in the sparse matrix, producing a tensor of shape
``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
or column (``dim=1``) dimension, producing a tensor of shape
``(input.shape[1],) + input.val.shape[1:]`` or
``(input.shape[0],) + input.val.shape[1:]``.
Returns
----------
torch.Tensor
Reduced tensor
Examples
----------
Case1: scalar-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.smax(A)
tensor(2)
>>> dglsp.smax(A, 0)
tensor([1, 0, 2])
>>> dglsp.smax(A, 1)
tensor([1, 2, 0, 0])
Case2: vector-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.smax(A)
tensor([2, 2])
>>> dglsp.smax(A, 1)
tensor([[1, 2],
[2, 2],
[0, 0],
[0, 0]])
"""
return torch.ops.dgl_sparse.smax(input.c_sparse_matrix, dim)
def smin(input: SparseMatrix, dim: Optional[int] = None):
"""Computes the minimum of non-zero values of the :attr:`input` sparse
matrix along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be reduced
does not have any non-zero value, the result will be 0.
Parameters
----------
input : SparseMatrix
The input sparse matrix
dim : int, optional
The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on both rows and columns simultaneously)
If :attr:`dim` is None, it reduces both the rows and the columns
in the sparse matrix, producing a tensor of shape
``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
or column (``dim=1``) dimension, producing a tensor of shape
``(input.shape[1],) + input.val.shape[1:]`` or
``(input.shape[0],) + input.val.shape[1:]``.
Returns
----------
torch.Tensor
Reduced tensor
Examples
----------
Case1: scalar-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.smin(A)
tensor(1)
>>> dglsp.smin(A, 0)
tensor([1, 0, 2])
>>> dglsp.smin(A, 1)
tensor([1, 1, 0, 0])
Case2: vector-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.smin(A)
tensor([1, 1])
>>> dglsp.smin(A, 0)
tensor([[1, 1],
[0, 0],
[2, 2]])
>>> dglsp.smin(A, 1)
tensor([[1, 2],
[2, 1],
[0, 0],
[0, 0]])
"""
return torch.ops.dgl_sparse.smin(input.c_sparse_matrix, dim)
def smean(input: SparseMatrix, dim: Optional[int] = None):
"""Computes the mean of non-zero values of the :attr:`input` sparse matrix
along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be reduced
does not have any non-zero value, the result will be 0.
Parameters
----------
input : SparseMatrix
The input sparse matrix
dim : int, optional
The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on both rows and columns simultaneously)
If :attr:`dim` is None, it reduces both the rows and the columns
in the sparse matrix, producing a tensor of shape
``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
or column (``dim=1``) dimension, producing a tensor of shape
``(input.shape[1],) + input.val.shape[1:]`` or
``(input.shape[0],) + input.val.shape[1:]``.
Returns
----------
torch.Tensor
Reduced tensor
Examples
----------
Case1: scalar-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([1., 1., 2.])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.smean(A)
tensor(1.3333)
>>> dglsp.smean(A, 0)
tensor([1., 0., 2.])
>>> dglsp.smean(A, 1)
tensor([1.0000, 1.5000, 0.0000, 0.0000])
Case2: vector-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([[1., 2.], [2., 1.], [2., 2.]])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.smean(A)
tensor([1.6667, 1.6667])
>>> dglsp.smean(A, 0)
tensor([[1.5000, 1.5000],
[0.0000, 0.0000],
[2.0000, 2.0000]])
>>> dglsp.smean(A, 1)
tensor([[1.0000, 2.0000],
[2.0000, 1.5000],
[0.0000, 0.0000],
[0.0000, 0.0000]])
"""
return torch.ops.dgl_sparse.smean(input.c_sparse_matrix, dim)
def sprod(input: SparseMatrix, dim: Optional[int] = None):
"""Computes the product of non-zero values of the :attr:`input` sparse
matrix along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be reduced
does not have any non-zero value, the result will be 0.
Parameters
----------
input : SparseMatrix
The input sparse matrix
dim : int, optional
The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on both rows and columns simultaneously)
If :attr:`dim` is None, it reduces both the rows and the columns
in the sparse matrix, producing a tensor of shape
``input.val.shape[1:]``. Otherwise, it reduces on the row (``dim=0``)
or column (``dim=1``) dimension, producing a tensor of shape
``(input.shape[1],) + input.val.shape[1:]`` or
``(input.shape[0],) + input.val.shape[1:]``.
Returns
----------
torch.Tensor
Reduced tensor
Examples
----------
Case1: scalar-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.sprod(A)
tensor(2)
>>> dglsp.sprod(A, 0)
tensor([1, 0, 2])
>>> dglsp.sprod(A, 1)
tensor([1, 2, 0, 0])
Case2: vector-valued sparse matrix
>>> indices = torch.tensor([[0, 1, 1], [0, 0, 2]])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.spmatrix(indices, val, shape=(4, 3))
>>> dglsp.sprod(A)
tensor([4, 4])
>>> dglsp.sprod(A, 0)
tensor([[2, 2],
[0, 0],
[2, 2]])
>>> dglsp.sprod(A, 1)
tensor([[1, 2],
[4, 2],
[0, 0],
[0, 0]])
"""
return torch.ops.dgl_sparse.sprod(input.c_sparse_matrix, dim)
SparseMatrix.reduce = reduce
SparseMatrix.sum = sum
SparseMatrix.smax = smax
SparseMatrix.smin = smin
SparseMatrix.smean = smean
SparseMatrix.sprod = sprod
python/dgl/sparse/sddmm.py
+104
View File
@@ -0,0 +1,104 @@
"""Sampled Dense-Dense Matrix Multiplication (SDDMM) operator module."""
import torch
from .sparse_matrix import SparseMatrix
__all__ = ["sddmm", "bsddmm"]
# pylint: disable=invalid-name
def sddmm(A: SparseMatrix, X1: torch.Tensor, X2: torch.Tensor) -> SparseMatrix:
r"""Sampled-Dense-Dense Matrix Multiplication (SDDMM).
``sddmm`` matrix-multiplies two dense matrices :attr:`X1` and :attr:`X2`,
then elementwise-multiplies the result with sparse matrix :attr:`A` at the
nonzero locations.
Mathematically ``sddmm`` is formulated as:
.. math::
out = (X1 @ X2) * A
In particular, :attr:`X1` and :attr:`X2` can be 1-D, then ``X1 @ X2``
becomes the out-product of the two vectors (which results in a matrix).
Parameters
----------
A : SparseMatrix
Sparse matrix of shape ``(L, N)``
X1 : torch.Tensor
Dense matrix of shape ``(L, M)`` or ``(L,)``
X2 : torch.Tensor
Dense matrix of shape ``(M, N)`` or ``(N,)``
Returns
-------
SparseMatrix
Sparse matrix of shape ``(L, N)``
Examples
--------
>>> indices = torch.tensor([[1, 1, 2], [2, 3, 3]])
>>> val = torch.arange(1, 4).float()
>>> A = dglsp.spmatrix(indices, val, (3, 4))
>>> X1 = torch.randn(3, 5)
>>> X2 = torch.randn(5, 4)
>>> dglsp.sddmm(A, X1, X2)
SparseMatrix(indices=tensor([[1, 1, 2],
[2, 3, 3]]),
values=tensor([-1.6585, -3.9714, -0.5406]),
shape=(3, 4), nnz=3)
"""
return SparseMatrix(torch.ops.dgl_sparse.sddmm(A.c_sparse_matrix, X1, X2))
# pylint: disable=invalid-name
def bsddmm(A: SparseMatrix, X1: torch.Tensor, X2: torch.Tensor) -> SparseMatrix:
r"""Sampled-Dense-Dense Matrix Multiplication (SDDMM) by batches.
``sddmm`` matrix-multiplies two dense matrices :attr:`X1` and :attr:`X2`,
then elementwise-multiplies the result with sparse matrix :attr:`A` at the
nonzero locations.
Mathematically ``sddmm`` is formulated as:
.. math::
out = (X1 @ X2) * A
The batch dimension is the last dimension for input dense matrices. In
particular, if the sparse matrix has scalar non-zero values, it will be
broadcasted for bsddmm.
Parameters
----------
A : SparseMatrix
Sparse matrix of shape ``(L, N)`` with scalar values or vector values of
length ``K``
X1 : Tensor
Dense matrix of shape ``(L, M, K)``
X2 : Tensor
Dense matrix of shape ``(M, N, K)``
Returns
-------
SparseMatrix
Sparse matrix of shape ``(L, N)`` with vector values of length ``K``
Examples
--------
>>> indices = torch.tensor([[1, 1, 2], [2, 3, 3]])
>>> val = torch.arange(1, 4).float()
>>> A = dglsp.spmatrix(indices, val, (3, 4))
>>> X1 = torch.arange(0, 3 * 5 * 2).view(3, 5, 2).float()
>>> X2 = torch.arange(0, 5 * 4 * 2).view(5, 4, 2).float()
>>> dglsp.bsddmm(A, X1, X2)
SparseMatrix(indices=tensor([[1, 1, 2],
[2, 3, 3]]),
values=tensor([[1560., 1735.],
[3400., 3770.],
[8400., 9105.]]),
shape=(3, 4), nnz=3, val_size=(2,))
"""
return sddmm(A, X1, X2)
python/dgl/sparse/softmax.py
+75
View File
@@ -0,0 +1,75 @@
"""Softmax op for SparseMatrix"""
# pylint: disable=invalid-name, W0622
import torch
from .sparse_matrix import SparseMatrix
__all__ = ["softmax"]
def softmax(input: SparseMatrix, dim: int = 1) -> SparseMatrix:
"""Applies softmax to the non-zero elements of the sparse matrix on the
dimension :attr:``dim``. dim = 0 or 1 indicates column-wise or row-wise
softmax respectively.
If :attr:`input.val` takes shape ``(nnz, D)``, then the output matrix
:attr:`output` and :attr:`output.val` take the same shape as :attr:`input`
and :attr:`input.val`. :attr:`output.val[:, i]` is calculated based on
:attr:`input.val[:, i]`.
Parameters
----------
input : SparseMatrix
The input sparse matrix
Returns
-------
SparseMatrix
The output sparse matrix
Examples
--------
Case1: row-wise softmax on matrix with values of shape (nnz)
>>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
>>> val = torch.tensor([0., 1., 2., 3.])
>>> A = dglsp.spmatrix(indices, val)
>>> dglsp.softmax(A)
SparseMatrix(indices=tensor([[0, 0, 1, 2],
[1, 2, 2, 0]]),
values=tensor([0.2689, 0.7311, 1.0000, 1.0000]),
shape=(3, 3), nnz=4)
Case2: row-wise softmax on matrix with values of shape (nnz, D)
>>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
>>> val = torch.tensor([[0., 7.], [1., 3.], [2., 2.], [3., 1.]])
>>> A = dglsp.spmatrix(indices, val)
>>> dglsp.softmax(A)
SparseMatrix(indices=tensor([[0, 0, 1, 2],
[1, 2, 2, 0]]),
values=tensor([[0.2689, 0.9820],
[0.7311, 0.0180],
[1.0000, 1.0000],
[1.0000, 1.0000]]),
shape=(3, 3), nnz=4, val_size=(2,))
Case3: column-wise softmax on matrix with values of shape (nnz)
>>> indices = torch.tensor([[0, 0, 1, 2], [1, 2, 2, 0]])
>>> val = torch.tensor([0., 1., 2., 3.])
>>> A = dglsp.spmatrix(indices, val)
>>> dglsp.softmax(A, 0)
SparseMatrix(indices=tensor([[0, 0, 1, 2],
[1, 2, 2, 0]]),
values=tensor([1.0000, 0.2689, 0.7311, 1.0000]),
shape=(3, 3), nnz=4)
"""
return SparseMatrix(
torch.ops.dgl_sparse.softmax(input.c_sparse_matrix, dim)
)
SparseMatrix.softmax = softmax
python/dgl/sparse/sparse_matrix.py
+1477
View File
File diff suppressed because it is too large Load Diff
python/dgl/sparse/unary_op.py
+61
View File
@@ -0,0 +1,61 @@
"""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
python/dgl/sparse/utils.py
+14
View File
@@ -0,0 +1,14 @@
"""Utilities for DGL sparse module."""
from numbers import Number
from typing import Union
import torch
def is_scalar(x):
"""Check if the input is a scalar."""
return isinstance(x, Number) or (torch.is_tensor(x) and x.dim() == 0)
# Scalar type annotation
Scalar = Union[Number, torch.Tensor]