Is there accumarray equivalent in pytorch?

Question

Let's assume I have a feature tensor [f1, f2, f3, f4].
I want to pool the features according to an arbitrary index tensor (e.g. [0,2,1,0]).
Then the results would be [f1+f4, f3, f2].

I found that accumarray is the function that I want, but it is on MATLAB.
Is there any similar one in PyTorch, while preserving gradient for learning?

Answer 1

torch_scatter.scatter_add works perfectly.