|
|
|
|
|
from typing import Optional |
|
|
|
import torch |
|
|
|
|
|
def broadcast(src: torch.Tensor, other: torch.Tensor, dim: int) -> torch.Tensor: |
|
if dim < 0: |
|
dim = other.dim() + dim |
|
if src.dim() == 1: |
|
for _ in range(0, dim): |
|
src = src.unsqueeze(0) |
|
for _ in range(src.dim(), other.dim()): |
|
src = src.unsqueeze(-1) |
|
src = src.expand(other.size()) |
|
return src |
|
|
|
|
|
def scatter_sum( |
|
src: torch.Tensor, |
|
index: torch.Tensor, |
|
dim: int = -1, |
|
out: Optional[torch.Tensor] = None, |
|
dim_size: Optional[int] = None, |
|
) -> torch.Tensor: |
|
index = broadcast(index, src, dim) |
|
if out is None: |
|
size = list(src.size()) |
|
if dim_size is not None: |
|
size[dim] = dim_size |
|
elif index.numel() == 0: |
|
size[dim] = 0 |
|
else: |
|
size[dim] = int(index.max()) + 1 |
|
out = torch.zeros(size, dtype=src.dtype, device=src.device) |
|
return out.scatter_add_(dim, index, src) |
|
else: |
|
return out.scatter_add_(dim, index, src) |
|
|
|
|
|
def scatter_add( |
|
src: torch.Tensor, |
|
index: torch.Tensor, |
|
dim: int = -1, |
|
out: Optional[torch.Tensor] = None, |
|
dim_size: Optional[int] = None, |
|
) -> torch.Tensor: |
|
return scatter_sum(src, index, dim, out, dim_size) |
|
|
|
|
|
def scatter_mul( |
|
src: torch.Tensor, |
|
index: torch.Tensor, |
|
dim: int = -1, |
|
out: Optional[torch.Tensor] = None, |
|
dim_size: Optional[int] = None, |
|
) -> torch.Tensor: |
|
return torch.ops.torch_scatter.scatter_mul(src, index, dim, out, dim_size) |
|
|
|
|
|
def scatter_mean( |
|
src: torch.Tensor, |
|
index: torch.Tensor, |
|
dim: int = -1, |
|
out: Optional[torch.Tensor] = None, |
|
dim_size: Optional[int] = None, |
|
) -> torch.Tensor: |
|
out = scatter_sum(src, index, dim, out, dim_size) |
|
dim_size = out.size(dim) |
|
|
|
index_dim = dim |
|
if index_dim < 0: |
|
index_dim = index_dim + src.dim() |
|
if index.dim() <= index_dim: |
|
index_dim = index.dim() - 1 |
|
|
|
ones = torch.ones(index.size(), dtype=src.dtype, device=src.device) |
|
count = scatter_sum(ones, index, index_dim, None, dim_size) |
|
count[count < 1] = 1 |
|
count = broadcast(count, out, dim) |
|
if out.is_floating_point(): |
|
out.true_divide_(count) |
|
else: |
|
out.div_(count, rounding_mode="floor") |
|
return out |
|
|
|
|
|
def scatter_min( |
|
src: torch.Tensor, |
|
index: torch.Tensor, |
|
dim: int = -1, |
|
out: Optional[torch.Tensor] = None, |
|
dim_size: Optional[int] = None, |
|
) -> tuple[torch.Tensor, torch.Tensor]: |
|
return torch.ops.torch_scatter.scatter_min(src, index, dim, out, dim_size) |
|
|
|
|
|
def scatter_max( |
|
src: torch.Tensor, |
|
index: torch.Tensor, |
|
dim: int = -1, |
|
out: Optional[torch.Tensor] = None, |
|
dim_size: Optional[int] = None, |
|
) -> tuple[torch.Tensor, torch.Tensor]: |
|
return torch.ops.torch_scatter.scatter_max(src, index, dim, out, dim_size) |
|
|
|
|
|
def scatter( |
|
src: torch.Tensor, |
|
index: torch.Tensor, |
|
dim: int = -1, |
|
out: Optional[torch.Tensor] = None, |
|
dim_size: Optional[int] = None, |
|
reduce: str = "sum", |
|
) -> torch.Tensor: |
|
r""" |
|
| |
|
|
|
.. image:: https://raw.githubusercontent.com/rusty1s/pytorch_scatter/ |
|
master/docs/source/_figures/add.svg?sanitize=true |
|
:align: center |
|
:width: 400px |
|
|
|
| |
|
|
|
Reduces all values from the :attr:`src` tensor into :attr:`out` at the |
|
indices specified in the :attr:`index` tensor along a given axis |
|
:attr:`dim`. |
|
For each value in :attr:`src`, its output index is specified by its index |
|
in :attr:`src` for dimensions outside of :attr:`dim` and by the |
|
corresponding value in :attr:`index` for dimension :attr:`dim`. |
|
The applied reduction is defined via the :attr:`reduce` argument. |
|
|
|
Formally, if :attr:`src` and :attr:`index` are :math:`n`-dimensional |
|
tensors with size :math:`(x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})` |
|
and :attr:`dim` = `i`, then :attr:`out` must be an :math:`n`-dimensional |
|
tensor with size :math:`(x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})`. |
|
Moreover, the values of :attr:`index` must be between :math:`0` and |
|
:math:`y - 1`, although no specific ordering of indices is required. |
|
The :attr:`index` tensor supports broadcasting in case its dimensions do |
|
not match with :attr:`src`. |
|
|
|
For one-dimensional tensors with :obj:`reduce="sum"`, the operation |
|
computes |
|
|
|
.. math:: |
|
\mathrm{out}_i = \mathrm{out}_i + \sum_j~\mathrm{src}_j |
|
|
|
where :math:`\sum_j` is over :math:`j` such that |
|
:math:`\mathrm{index}_j = i`. |
|
|
|
.. note:: |
|
|
|
This operation is implemented via atomic operations on the GPU and is |
|
therefore **non-deterministic** since the order of parallel operations |
|
to the same value is undetermined. |
|
For floating-point variables, this results in a source of variance in |
|
the result. |
|
|
|
:param src: The source tensor. |
|
:param index: The indices of elements to scatter. |
|
:param dim: The axis along which to index. (default: :obj:`-1`) |
|
:param out: The destination tensor. |
|
:param dim_size: If :attr:`out` is not given, automatically create output |
|
with size :attr:`dim_size` at dimension :attr:`dim`. |
|
If :attr:`dim_size` is not given, a minimal sized output tensor |
|
according to :obj:`index.max() + 1` is returned. |
|
:param reduce: The reduce operation (:obj:`"sum"`, :obj:`"mul"`, |
|
:obj:`"mean"`, :obj:`"min"` or :obj:`"max"`). (default: :obj:`"sum"`) |
|
|
|
:rtype: :class:`Tensor` |
|
|
|
.. code-block:: python |
|
|
|
from torch_scatter import scatter |
|
|
|
src = torch.randn(10, 6, 64) |
|
index = torch.tensor([0, 1, 0, 1, 2, 1]) |
|
|
|
# Broadcasting in the first and last dim. |
|
out = scatter(src, index, dim=1, reduce="sum") |
|
|
|
print(out.size()) |
|
|
|
.. code-block:: |
|
|
|
torch.Size([10, 3, 64]) |
|
""" |
|
if reduce == "sum" or reduce == "add": |
|
return scatter_sum(src, index, dim, out, dim_size) |
|
if reduce == "mul": |
|
return scatter_mul(src, index, dim, out, dim_size) |
|
elif reduce == "mean": |
|
return scatter_mean(src, index, dim, out, dim_size) |
|
elif reduce == "min": |
|
return scatter_min(src, index, dim, out, dim_size)[0] |
|
elif reduce == "max": |
|
return scatter_max(src, index, dim, out, dim_size)[0] |
|
else: |
|
raise ValueError |
|
|
