Scatter Mean

torch_scatter.scatter_mean(src, index, dim=-1, out=None, dim_size=None, fill_value=0)[source]

https://raw.githubusercontent.com/rusty1s/pytorch_scatter/master/docs/source/_figures/mean.svg?sanitize=true

Averages all values from the src tensor into out at the indices specified in the index tensor along a given axis dim.If multiple indices reference the same location, their contributions average (cf. scatter_add()).

For one-dimensional tensors, the operation computes

\[\mathrm{out}_i = \mathrm{out}_i + \frac{1}{N_i} \cdot \sum_j \mathrm{src}_j\]

where \(\sum_j\) is over \(j\) such that \(\mathrm{index}_j = i\). \(N_i\) indicates the number of indices referencing \(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_mean

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_mean(src, index, out=out)

print(out)
tensor([[0.0000, 0.0000, 4.0000, 3.0000, 1.5000, 0.0000],
        [1.0000, 4.0000, 2.0000, 0.0000, 0.0000, 0.0000]])