How can I compute the mean of values selected from a vector A from an indexing vector B?
I have a vector of values, for example:

import torch
v = torch.rand(6)


tensor([0.0811, 0.9658, 0.1901, 0.0872, 0.8895, 0.9647])


and an index to select values from v:

index = torch.tensor([0, 1, 0, 2, 0, 2])


I want to produce a vector mean which would compute the mean values of v grouped by indexes from index.

In this example:

mean = torch.tensor([(0.0811 + 0.1901 + 0.8895) / 3, 0.9658, (0.0872 + 0.9647) / 2, 0, 0, 0])


tensor([0.3869, 0.9658, 0.5260, 0.0000, 0.0000, 0.0000])


One possible solution using a combination of torch.bincount and Tensor.index_add():

v = torch.tensor([0.0811, 0.9658, 0.1901, 0.0872, 0.8895, 0.9647])
index = torch.tensor([0, 1, 0, 2, 0, 2])


bincount() gets the total for each index use in index:

bincount = torch.bincount(index, minlength=6)
# --&gt; tensor([3, 1, 2, 0, 0, 0])


index_add() adds from v in the order given by index:

numerator = torch.zeros(6)
numerator = numerator.index_add(0, index, v)
# --&gt; tensor([1.1607, 0.9658, 1.0520, 0.0000, 0.0000, 0.0000])


Replace zeros with 1.0 in bincount to prevent division by 0 
and convert from int to float:

div = bincount.float()
div[bincount == 0] = 1.0
# --&gt; tensor([3., 1., 2., 1., 1., 1.])

mean = num/div
# --&gt; tensor([0.3869, 0.9658, 0.5260, 0.0000, 0.0000, 0.0000])

