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? Torch.gather seemed slightly similar but wasn’t able to do the job.Thanks!