Skip to content

GitLab

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

[Sparse] Polish reduction.py. (#5180) 



* polish

* example

* revert_example

* pylint

* revert_input

* order

* add input as parameter

* example
Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent b9c8d458
Hide whitespace changes
Inline Side-by-side
Showing with 103 additions and 92 deletions
python/dgl/sparse/reduction.py
View file @ 84aae086
"""DGL sparse matrix reduce operators""" """DGL sparse matrix reduce operators"""
# pylint: disable=W0622
from typing import Optional from typing import Optional
import torch import torch
...@@ -6,30 +8,30 @@ import torch ...@@ -6,30 +8,30 @@ import torch
from .sparse_matrix import SparseMatrix from .sparse_matrix import SparseMatrix
def reduce(A: SparseMatrix, dim: Optional[int] = None, rtype: str = "sum"): def reduce(input: SparseMatrix, dim: Optional[int] = None, rtype: str = "sum"):
"""Compute the reduction of non-zero values in sparse matrix A along """Computes the reduction of non-zero values of the ``input`` sparse matrix
the given dimension :attr:`dim`. along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be The reduction does not count zero elements. If the row or column to be
reduced does not have any non-zero value, the result will be 0. reduced does not have any non-zero elements, the result will be 0.
Parameters Parameters
---------- ----------
A : SparseMatrix input : SparseMatrix
Sparse matrix The input sparse matrix
dim : int, optional dim : int, optional
The dimension to reduce. Must be either 0 (by rows) or 1 (by columns) The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on all non-zero entries). or None (on all non-zero entries)
If :attr:`dim` is None, it reduces all the elements in the sparse If :attr:`dim` is None, it reduces all the elements in the sparse
matrix. Otherwise, it reduces on the row (``dim=0``) or column matrix. Otherwise, it reduces on the row (``dim=0``) or column
(``dim=1``) dimension, producing a tensor of shape (``dim=1``) dimension, producing a tensor of shape
``(A.shape[1],) + A.val.shape[1:]`` or ``(input.shape[1],) + input.val.shape[1:]`` or
``(A.shape[0],) + A.val.shape[1:]``. ``(input.shape[0],) + input.val.shape[1:]``.
rtype: str, optional rtype: str, optional
Reduction type, one of ``['sum', 'smin', 'smax', 'smean', 'sprod']``, Reduction type, one of ``['sum', 'smin', 'smax', 'smean', 'sprod']``,
representing taking the sum, minimum, maximum, mean, and product of the representing taking the sum, minimum, maximum, mean, and product of the
non-zero entries. non-zero elements
Returns Returns
---------- ----------
...@@ -43,56 +45,57 @@ def reduce(A: SparseMatrix, dim: Optional[int] = None, rtype: str = "sum"): ...@@ -43,56 +45,57 @@ def reduce(A: SparseMatrix, dim: Optional[int] = None, rtype: str = "sum"):
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([1, 1, 2]) >>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.reduce(rtype='sum')) >>> print(dglsp.reduce(A, rtype='sum'))
tensor(4) tensor(4)
>>> print(A.reduce(0, 'sum')) >>> print(dglsp.reduce(A, 0, 'sum'))
tensor([2, 0, 2]) tensor([2, 0, 2])
>>> print(A.reduce(1, 'sum')) >>> print(dglsp.reduce(A, 1, 'sum'))
tensor([1, 3, 0, 0]) tensor([1, 3, 0, 0])
>>> print(A.reduce(0, 'smax')) >>> print(dglsp.reduce(A, 0, 'smax'))
tensor([1, 0, 2]) tensor([1, 0, 2])
>>> print(A.reduce(1, 'smin')) >>> print(dglsp.reduce(A, 1, 'smin'))
tensor([1, 1, 0, 0]) tensor([1, 1, 0, 0])
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([[1., 2.], [2., 1.], [2., 2.]]) >>> val = torch.tensor([[1., 2.], [2., 1.], [2., 2.]])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.reduce(rtype='sum')) >>> print(dglsp.reduce(A, rtype='sum'))
tensor([5., 5.]) tensor([5., 5.])
>>> print(A.reduce(0, 'sum')) >>> print(dglsp.reduce(A, 0, 'sum'))
tensor([[3., 3.], tensor([[3., 3.],
[0., 0.], [0., 0.],
[2., 2.]]) [2., 2.]])
>>> print(A.reduce(1, 'smin')) >>> print(dglsp.reduce(A, 1, 'smin'))
tensor([[1., 2.], tensor([[1., 2.],
[2., 1.], [2., 1.],
[0., 0.], [0., 0.],
[0., 0.]]) [0., 0.]])
>>> print(A.reduce(0, 'smean')) >>> print(dglsp.reduce(A, 0, 'smean'))
tensor([[1.5000, 1.5000], tensor([[1.5000, 1.5000],
[0.0000, 0.0000], [0.0000, 0.0000],
[2.0000, 2.0000]]) [2.0000, 2.0000]])
""" """
return torch.ops.dgl_sparse.reduce(A.c_sparse_matrix, rtype, dim) return torch.ops.dgl_sparse.reduce(input.c_sparse_matrix, rtype, dim)
# pylint: disable=W0622 def sum(input: SparseMatrix, dim: Optional[int] = None):
def sum(A: SparseMatrix, dim: Optional[int] = None): """Computes the sum of non-zero values of the ``input`` sparse matrix along
"""Compute the sum of non-zero values in sparse matrix A along
the given dimension :attr:`dim`. the given dimension :attr:`dim`.
Parameters Parameters
---------- ----------
input : SparseMatrix
The input sparse matrix
dim : int, optional dim : int, optional
The dimension to reduce. Must be either 0 (by rows) or 1 (by columns) The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on all non-zero entries). or None (on all non-zero entries)
If :attr:`dim` is None, it reduces all the elements in the sparse If :attr:`dim` is None, it reduces all the elements in the sparse
matrix. Otherwise, it reduces on the row (``dim=0``) or column matrix. Otherwise, it reduces on the row (``dim=0``) or column
(``dim=1``) dimension, producing a tensor of shape (``dim=1``) dimension, producing a tensor of shape
``(A.shape[1],) + A.val.shape[1:]`` or ``(input.shape[1],) + input.val.shape[1:]`` or
``(A.shape[0],) + A.val.shape[1:]``. ``(input.shape[0],) + input.val.shape[1:]``.
Returns Returns
---------- ----------
...@@ -106,45 +109,47 @@ def sum(A: SparseMatrix, dim: Optional[int] = None): ...@@ -106,45 +109,47 @@ def sum(A: SparseMatrix, dim: Optional[int] = None):
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([1, 1, 2]) >>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.sum()) >>> print(dglsp.sum(A))
tensor(4) tensor(4)
>>> print(A.sum(0)) >>> print(dglsp.sum(A, 0))
tensor([2, 0, 2]) tensor([2, 0, 2])
>>> print(A.sum(1)) >>> print(dglsp.sum(A, 1))
tensor([1, 3, 0, 0]) tensor([1, 3, 0, 0])
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]]) >>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.sum()) >>> print(dglsp.sum(A))
tensor([5, 5]) tensor([5, 5])
>>> print(A.sum(0)) >>> print(dglsp.sum(A, 0))
tensor([[3, 3], tensor([[3, 3],
[0, 0], [0, 0],
[2, 2]]) [2, 2]])
""" """
return torch.ops.dgl_sparse.sum(A.c_sparse_matrix, dim) return torch.ops.dgl_sparse.sum(input.c_sparse_matrix, dim)
def smax(A: SparseMatrix, dim: Optional[int] = None): def smax(input: SparseMatrix, dim: Optional[int] = None):
"""Compute the maximum of non-zero values in sparse matrix A along """Computes the maximum of non-zero values of the ``input`` sparse matrix
the given dimension :attr:`dim`. along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be 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. reduced does not have any non-zero value, the result will be 0.
Parameters Parameters
---------- ----------
input : SparseMatrix
The input sparse matrix
dim : int, optional dim : int, optional
The dimension to reduce. Must be either 0 (by rows) or 1 (by columns) The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on all non-zero entries). or None (on all non-zero entries)
If :attr:`dim` is None, it reduces all the elements in the sparse If :attr:`dim` is None, it reduces all the elements in the sparse
matrix. Otherwise, it reduces on the row (``dim=0``) or column matrix. Otherwise, it reduces on the row (``dim=0``) or column
(``dim=1``) dimension, producing a tensor of shape (``dim=1``) dimension, producing a tensor of shape
``(A.shape[1],) + A.val.shape[1:]`` or ``(input.shape[1],) + input.val.shape[1:]`` or
``(A.shape[0],) + A.val.shape[1:]``. ``(input.shape[0],) + input.val.shape[1:]``.
Returns Returns
---------- ----------
...@@ -158,87 +163,89 @@ def smax(A: SparseMatrix, dim: Optional[int] = None): ...@@ -158,87 +163,89 @@ def smax(A: SparseMatrix, dim: Optional[int] = None):
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([1, 1, 2]) >>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.smax()) >>> print(dglsp.smax(A))
tensor(2) tensor(2)
>>> print(A.smax(0)) >>> print(dglsp.smax(A, 0))
tensor([1, 0, 2]) tensor([1, 0, 2])
>>> print(A.smax(1)) >>> print(dglsp.smax(A, 1))
tensor([1, 2, 0, 0]) tensor([1, 2, 0, 0])
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]]) >>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.smax()) >>> print(dglsp.smax(A))
tensor([2, 2]) tensor([2, 2])
>>> print(A.smax(1)) >>> print(dglsp.smax(A, 1))
tensor([[1, 2], tensor([[1, 2],
[2, 2], [2, 2],
[0, 0], [0, 0],
[0, 0]]) [0, 0]])
""" """
return torch.ops.dgl_sparse.smax(A.c_sparse_matrix, dim) return torch.ops.dgl_sparse.smax(input.c_sparse_matrix, dim)
def smin(A: SparseMatrix, dim: Optional[int] = None): def smin(input: SparseMatrix, dim: Optional[int] = None):
"""Compute the minimum of non-zero values in sparse matrix A along """Computes the minimum of non-zero values of the ``input`` sparse matrix
the given dimension :attr:`dim`. along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be reduced 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. does not have any non-zero value, the result will be 0.
Parameters Parameters
---------- ----------
input : SparseMatrix
The input sparse matrix
dim : int, optional dim : int, optional
The dimension to reduce. Must be either 0 (by rows) or 1 (by columns) The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on all non-zero entries). or None (on all non-zero entries)
If :attr:`dim` is None, it reduces all the elements in the sparse If :attr:`dim` is None, it reduces all the elements in the sparse
matrix. Otherwise, it reduces on the row (``dim=0``) or column matrix. Otherwise, it reduces on the row (``dim=0``) or column
(``dim=1``) dimension, producing a tensor of shape (``dim=1``) dimension, producing a tensor of shape
``(A.shape[1],) + A.val.shape[1:]`` or ``(input.shape[1],) + input.val.shape[1:]`` or
``(A.shape[0],) + A.val.shape[1:]``. ``(input.shape[0],) + input.val.shape[1:]``.
Returns Returns
---------- ----------
Tensor Tensor
Reduced tensor Reduced tensor
Example Examples
---------- ----------
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([1, 1, 2]) >>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.smin()) >>> print(dglsp.smin(A))
tensor(1) tensor(1)
>>> print(A.smin(0)) >>> print(dglsp.smin(A, 0))
tensor([1, 0, 2]) tensor([1, 0, 2])
>>> print(A.smin(1)) >>> print(dglsp.smin(A, 1))
tensor([1, 1, 0, 0]) tensor([1, 1, 0, 0])
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]]) >>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.smin()) >>> print(dglsp.smin(A))
tensor([1, 1]) tensor([1, 1])
>>> print(A.smin(0)) >>> print(dglsp.smin(A, 0))
tensor([[1, 1], tensor([[1, 1],
[0, 0], [0, 0],
[2, 2]]) [2, 2]])
>>> print(A.smin(1)) >>> print(dglsp.smin(A, 1))
tensor([[1, 2], tensor([[1, 2],
[2, 1], [2, 1],
[0, 0], [0, 0],
[0, 0]]) [0, 0]])
""" """
return torch.ops.dgl_sparse.smin(A.c_sparse_matrix, dim) return torch.ops.dgl_sparse.smin(input.c_sparse_matrix, dim)
def smean(A: SparseMatrix, dim: Optional[int] = None): def smean(input: SparseMatrix, dim: Optional[int] = None):
"""Compute the mean of non-zero values in sparse matrix A along """Computes the mean of non-zero values of the ``input`` sparse matrix along
the given dimension :attr:`dim`. the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be reduced The reduction does not count zero values. If the row or column to be reduced
...@@ -246,109 +253,113 @@ def smean(A: SparseMatrix, dim: Optional[int] = None): ...@@ -246,109 +253,113 @@ def smean(A: SparseMatrix, dim: Optional[int] = None):
Parameters Parameters
---------- ----------
input : SparseMatrix
The input sparse matrix
dim : int, optional dim : int, optional
The dimension to reduce. Must be either 0 (by rows) or 1 (by columns) The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on all non-zero entries). or None (on all non-zero entries)
If :attr:`dim` is None, it reduces all the elements in the sparse If :attr:`dim` is None, it reduces all the elements in the sparse
matrix. Otherwise, it reduces on the row (``dim=0``) or column matrix. Otherwise, it reduces on the row (``dim=0``) or column
(``dim=1``) dimension, producing a tensor of shape (``dim=1``) dimension, producing a tensor of shape
``(A.shape[1],) + A.val.shape[1:]`` or ``(input.shape[1],) + input.val.shape[1:]`` or
``(A.shape[0],) + A.val.shape[1:]``. ``(input.shape[0],) + input.val.shape[1:]``.
Returns Returns
---------- ----------
Tensor Tensor
Reduced tensor Reduced tensor
Example Examples
---------- ----------
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([1., 1., 2.]) >>> val = torch.tensor([1., 1., 2.])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.smean()) >>> print(dglsp.smean(A))
tensor(1.3333) tensor(1.3333)
>>> print(A.smean(0)) >>> print(dglsp.smean(A, 0))
tensor([1., 0., 2.]) tensor([1., 0., 2.])
>>> print(A.smean(1)) >>> print(dglsp.smean(A, 1))
tensor([1.0000, 1.5000, 0.0000, 0.0000]) tensor([1.0000, 1.5000, 0.0000, 0.0000])
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([[1., 2.], [2., 1.], [2., 2.]]) >>> val = torch.tensor([[1., 2.], [2., 1.], [2., 2.]])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.smean()) >>> print(dglsp.smean(A))
tensor([1.6667, 1.6667]) tensor([1.6667, 1.6667])
>>> print(A.smean(0)) >>> print(dglsp.smean(A, 0))
tensor([[1.5000, 1.5000], tensor([[1.5000, 1.5000],
[0.0000, 0.0000], [0.0000, 0.0000],
[2.0000, 2.0000]]) [2.0000, 2.0000]])
>>> print(A.smean(1)) >>> print(dglsp.smean(A, 1))
tensor([[1.0000, 2.0000], tensor([[1.0000, 2.0000],
[2.0000, 1.5000], [2.0000, 1.5000],
[0.0000, 0.0000], [0.0000, 0.0000],
[0.0000, 0.0000]]) [0.0000, 0.0000]])
""" """
return torch.ops.dgl_sparse.smean(A.c_sparse_matrix, dim) return torch.ops.dgl_sparse.smean(input.c_sparse_matrix, dim)
def sprod(A: SparseMatrix, dim: Optional[int] = None): def sprod(input: SparseMatrix, dim: Optional[int] = None):
"""Compute the product of non-zero values in sparse matrix A along """Computes the product of non-zero values of the ``input`` sparse matrix
the given dimension :attr:`dim`. along the given dimension :attr:`dim`.
The reduction does not count zero values. If the row or column to be reduced 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. does not have any non-zero value, the result will be 0.
Parameters Parameters
---------- ----------
input : SparseMatrix
The input sparse matrix
dim : int, optional dim : int, optional
The dimension to reduce. Must be either 0 (by rows) or 1 (by columns) The dimension to reduce, must be either 0 (by rows) or 1 (by columns)
or None (on all non-zero entries). or None (on all non-zero entries)
If :attr:`dim` is None, it reduces all the elements in the sparse If :attr:`dim` is None, it reduces all the elements in the sparse
matrix. Otherwise, it reduces on the row (``dim=0``) or column matrix. Otherwise, it reduces on the row (``dim=0``) or column
(``dim=1``) dimension, producing a tensor of shape (``dim=1``) dimension, producing a tensor of shape
``(A.shape[1],) + A.val.shape[1:]`` or ``(input.shape[1],) + input.val.shape[1:]`` or
``(A.shape[0],) + A.val.shape[1:]``. ``(input.shape[0],) + input.val.shape[1:]``.
Returns Returns
---------- ----------
Tensor Tensor
Reduced tensor Reduced tensor
Example Examples
---------- ----------
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([1, 1, 2]) >>> val = torch.tensor([1, 1, 2])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.sprod()) >>> print(dglsp.sprod(A))
tensor(2) tensor(2)
>>> print(A.sprod(0)) >>> print(dglsp.sprod(A, 0))
tensor([1, 0, 2]) tensor([1, 0, 2])
>>> print(A.sprod(1)) >>> print(dglsp.sprod(A, 1))
tensor([1, 2, 0, 0]) tensor([1, 2, 0, 0])
>>> row = torch.tensor([0, 1, 1]) >>> row = torch.tensor([0, 1, 1])
>>> col = torch.tensor([0, 0, 2]) >>> col = torch.tensor([0, 0, 2])
>>> val = torch.tensor([[1, 2], [2, 1], [2, 2]]) >>> val = torch.tensor([[1, 2], [2, 1], [2, 2]])
>>> A = dglsp.from_coo(row, col, val, shape=(4, 3)) >>> A = dglsp.from_coo(row, col, val, shape=(4, 3))
>>> print(A.sprod()) >>> print(dglsp.sprod(A))
tensor([4, 4]) tensor([4, 4])
>>> print(A.sprod(0)) >>> print(dglsp.sprod(A, 0))
tensor([[2, 2], tensor([[2, 2],
[0, 0], [0, 0],
[2, 2]]) [2, 2]])
>>> print(A.sprod(1)) >>> print(dglsp.sprod(A, 1))
tensor([[1, 2], tensor([[1, 2],
[4, 2], [4, 2],
[0, 0], [0, 0],
[0, 0]]) [0, 0]])
""" """
return torch.ops.dgl_sparse.sprod(A.c_sparse_matrix, dim) return torch.ops.dgl_sparse.sprod(input.c_sparse_matrix, dim)
SparseMatrix.reduce = reduce SparseMatrix.reduce = reduce
......
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