CODEWITHSUNDEEP
pythonmachine-learningpytorch
mshuaibi
mshuaibi

Reputation: 119

Tensor reduction based off index vector

As an example, I have 2 tensors: A = [1;2;3;4;5;6;7] and B = [2;3;2]. The idea is that I want to reduce A based off B - such that B's values represent how to sum A's values- such that B = [2;3;2] means the reduced A shall be the sum of the first 2 values, next 3, and last 2: A' = [(1+2);(3+4+5);(6+7)]. It is apparent that the sum of B shall always be equal to the length of A. I'm trying to do this as efficiently as possible - preferably specific functions or matrix operations contained within pytorch/python. Thanks!

Upvotes: 0

Views: 602

Answers (1)

David Ng
David Ng

Reputation: 1698

Here is the solution.

  • First, we create an array of indices B_idx with the same size of A.
  • Then, accumulate (add) all elements in A based on the indices B_idx using index_add_.
A = torch.arange(1, 8) 
B = torch.tensor([2, 3, 2])

B_idx = [idx.repeat(times) for idx, times in zip(torch.arange(len(B)), B)]
B_idx = torch.cat(B_idx) # tensor([0, 0, 1, 1, 1, 2, 2])

A_sum = torch.zeros_like(B)
A_sum.index_add_(dim=0, index=B_idx, source=A)
print(A_sum) # tensor([ 3, 12, 13])

Upvotes: 1

Related Questions