Batched index_add_

Hi,
I wanted to have index_add_ operation applied on a batch.
For instance:
A = torch.FloatTensor([[[1, 2, 3], [1, 2, 3]], [[4, 5, 6], [4, 5, 6]], [[7, 8, 9], [7,8, 9]]])
ind = torch.LongTensor([[0, 1], [1, 2], [0, 2]])
to_add = torch.LongTensor([[[0,1],[0,1]],[[2,3],[2,3]],[[4,5],[4,5]]])

I would have liked a batched index_add_ operation
A.index_add_(2, ind, to_add)
and the expected output as:
output = torch.FloatTensor([[[1, 3, 3], [1, 3, 3]], [[6, 8, 6], [6, 8, 6]], [[11, 13, 9], [11, 13, 9]]])

Currently I am achieving the same using a loop through each tensor in batch and using index_add_ on each tensor. I would like to avoid this loop. Please let me know if there is a more efficient way of doing the same.

Thanks.

Hi,

I just met the same problem as you, and I found torch.scatter_add_() is helpful! You can have a try.

Btw, I found the methods named gather and scatter in PyTorch can solve this kind of problems well, which can be seemed as a “universal” version of methods named index*.