Scatter Add¶
-
torch_scatter.
scatter_add
(src, index, dim=-1, out=None, dim_size=None, fill_value=0)[source]¶ Sums all values from the
src
tensor intoout
at the indices specified in theindex
tensor along a given axisdim
. For each value insrc
, its output index is specified by its index ininput
for dimensions outside ofdim
and by the corresponding value inindex
for dimensiondim
. If multiple indices reference the same location, their contributions add.Formally, if
src
andindex
are n-dimensional tensors with size \((x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})\) anddim
= i, thenout
must be an n-dimensional tensor with size \((x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})\). Moreover, the values ofindex
must be between 0 and out.size(dim) - 1.For one-dimensional tensors, the operation computes
\[\mathrm{out}_i = \mathrm{out}_i + \sum_j \mathrm{src}_j\]where \(\sum_j\) is over \(j\) such that \(\mathrm{index}_j = i\).
Parameters: - src (Tensor) – The source tensor.
- index (LongTensor) – The indices of elements to scatter.
- dim (int, optional) – The axis along which to index.
(default:
-1
) - out (Tensor, optional) – The destination tensor. (default:
None
) - dim_size (int, optional) – If
out
is not given, automatically create output with sizedim_size
at dimensiondim
. Ifdim_size
is not given, a minimal sized output tensor is returned. (default:None
) - fill_value (int, optional) – If
out
is not given, automatically fill output tensor withfill_value
. (default:0
)
Return type: Tensor
from torch_scatter import scatter_add src = torch.Tensor([[2, 0, 1, 4, 3], [0, 2, 1, 3, 4]]) index = torch.tensor([[4, 5, 4, 2, 3], [0, 0, 2, 2, 1]]) out = src.new_zeros((2, 6)) out = scatter_add(src, index, out=out) print(out)
tensor([[0., 0., 4., 3., 3., 0.], [2., 4., 4., 0., 0., 0.]])