How to do atomic add on slice with duplicate indices?

Hi everyone! I’m currently working on my project, in which I need to select entries of a tensor based on an index tensor and add values to the selected entries. However, I found out that an inplace add may not always work because of data races. For example, I have an index tensor which may have some duplicate indices, like below:

tensor = torch.tensor([1, 2, 3])
data = torch.tensor([4, 5, 6])
indices = torch.tensor([0, 0, 1], dtype = torch.long)
tensor[indices] += data

What I expect is tensor: [10, 8, 3] because I want 4 and 5 to be added to tensor[0], but what I get is tensor: [6, 8, 3], which means only 5 is added to tensor[0]. I also encountered this problem when I tried this with 2D and 3D tensors.

I searched this issue online but found nothing useful, can you please help me with this? Thanks!

Update:

This 1D case can be resolved by torch.scatter_add, but it seems that scatter_add cannot deal with the below case

tensor = torch.tensor([[1, 2, 3],
                       [11, 22, 33]])
data = torch.tensor([[4, 5, 6],
                     [7, 8, 9],
                     [10, 11, 12]])
indices_x = torch.ones_like(data)
indices_y = torch.zeros_like(data)
tensor[indices_x, indices_y] += data  # expect to add all values of data to tensor[1, 0]
print(tensor)

It might be easiest to linearize the indices (flattened = tensor.view(-1) and indices = indices_x * tensor.dim(1) + indices_y). Personally, I’d probably use index_add btw.
It’s a bit more clumsy to write, but in the end a multi-dimensional index addition wowuld do the same under the hood (but probably on the fly without materializing the index tensor).

Best regards

Thomas

Thanks a lot! That works, although it looks not so nice.

Glad it works! Maybe if you use a cuter font, e.g. Roses Bolero by Authentype:

image

More seriously, yes, be sure to put in a bit of commentary as what these things do.

Best regards

Thomas