Scatter Add

torch_scatter.scatter_add(src, index, dim=-1, out=None, dim_size=None, fill_value=0)

Sums all values from the src tensor into out at the indices specified in the index tensor along a given axis dim. For each value in src, its output index is specified by its index in input for dimensions outside of dim and by the corresponding value in index for dimension dim. If multiple indices reference the same location, their contributions add.

Formally, if src and index are n-dimensional tensors with size \((x_0, ..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})\) and dim = i, then out must be an n-dimensional tensor with size \((x_0, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})\). Moreover, the values of index must be between 0 and out.size(dim) - 1. Both src and index are broadcasted in case their dimensions do not match.

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 size dim_size at dimension dim. If dim_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 with fill_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.]])