nn.scatter_add

nn.scatter_add(x: torch.Tensor, idx_i: torch.Tensor, dim_size: int, dim: int = 0) → torch.Tensor[source]

Sum over values with the same indices.

Parameters:

  • x – input values

  • idx_i – index of center atom i

  • dim_size – size of the dimension after reduction

  • dim – the dimension to reduce

Returns:

reduced input