How to sum based off index vector
I have 3 vectors - a sum vector, a contribution vector, and a value vector. I want to sum the value vectors according to their contribution vector and place them in their corresponding index in the sum vector. An example is:

A = [0;0] (sum vector), B = [0,0,1,1] (contribution vector) C=[20,30,40,10] (value vector)

Output:
A = [20+30;40+10]

Such that the B vector is the same length as C and their corresponding index tell us what position in A to be added to.

I am able to achieve this by a for loop as such:

for index,value in enumerate(C):
    A[B[index]]+=value


However, as this will be part of my NN model forward loop it will cause significant performance issue. Specifically I was looking for a vector/matrix sorting approach that will be more efficient. In the example above, something that worked efficiently for me was:

A=torch.zeros(2,1)
C=C.reshape(2,2)
sum=torch.sum(C,1).reshape(2,1)
A += sum


However, I run into issues as it is not always the case that the indexes of A have the same contribution. For example - the case such that B = [0,0,0,1,1] and C=[20,30,40,10,50]. Is there a function or a strategic way to do this for general cases? Thanks!

You are looking for index_add_()

A.index_add_(0, B, C)


Note that B should be of type torch.long (it is an index vector), and C should be of type torch.float, same as the type of A.
Moreover, you can use the first dim argument to do this summation along different dimensions in case A and C are multi-dimensional tensors.
