Element-wise matrix addition of two differently shaped matrices (tensors) based on index

1

Hello all,

I am curious if you guys have a better (i.e., faster, computationally more efficient) solution than me.

Situation:
I have two tensors with size A=[bs, n, k] and B=[bs, m, k], where bs: batch size, n: number of low-level labels, m: number of high-level categories, and k: hidden dimension. I want to add the elements of B to specific elements/vectors of A based on a known index. I know exactly which element of B has to be added to one or more of the elements of A. Please see a working example of how I approach this problem (disregard the batch size as it is not important for the problem) - please note that I perform these operations on multiple GPUs:

import torch

A = torch.tensor([[0, 0, 0,], [1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]])
B = torch.tensor([[1, 1, 1], [2, 2, 2]])

indices = [0, 1, 1, 0, 0]

for i, index in enumerate(indices):
A[i] += B[index]

print(A)

Result:
tensor([[1, 1, 1],
[3, 3, 3],
[4, 4, 4],
[4, 4, 4],
[5, 5, 5]])

Is there a more efficient way to this?

Thank you,
Christoph

2

Direct indexing should work:

A_ = torch.tensor([[0, 0, 0,], [1, 1, 1], [2, 2, 2], [3, 3, 3], [4, 4, 4]])
A_[torch.arange(len(indices))] += B[indices]

print((A_ == A).all())
# tensor(True)
3

That will work, thank you!