Torch - How to calculate average of tensors with the same indexes
Suppose having two matrices: X(m, n) and index matrix I(m, 1). Every item in index matrix I_k represents the index of the kth element X_k in X.
And suppose the index is in the range of [0, 1, 2, ..., j-1]
I would like to calculate the average of tensors in X with the same index i and return a result matrix R(j, n).
For example,
X = [[1, 1, 1],
     [2, 2, 2],
     [3, 3, 3]]

I = [0, 0, 1]

The result matrix should be:
R = [[torch.mean(([1, 1, 1], [2, 2, 2]))],
     [torch.mean(([3, 3, 3]))]

which equals to:
R = [[1.5, 1.5, 1.5],
     [3, 3, 3]]

My current solution is to traverse through m, stack the tensors with the same index and perform torch.mean.
Is there a way avoiding traversing through m? It seems not elegant and rather time-consuming.

ret = torch.empty_like(X)
ret.scatter_reduce_(0, I.unsqueeze(-1).expand_as(X), X, &quot;mean&quot;, include_self=False)

should do what you want.
Now, note that this is a fairly new method so it may not be particularly performant. If you bump into an issue with this method, you may be better off running scatter_add_ on the tensor X and a tensor of ones and then divide.
If you want to also have a smaller tensor as output, you may want to figure out how many indices and with that infer the size of out.
