Summing over specific indices PyTorch (similar to scatter_add)
I have some matrix where rows belong to some label, unordered. I want to sum all rows for each label.
Here is how it can be done with a loop:
labels = torch.tensor([0, 1, 0])
x = torch.tensor([[1, 2, 3],[4, 5, 6],[7, 8, 9]])
torch.stack([torch.sum(x[labels == i], dim=0) for i in torch.unique(labels)])

desired output:
tensor([[ 8, 10, 12],
        [ 4,  5,  6]])

EDIT: Just to make it clear, I have the labels tensor, I know which labels repeat, I am interested in computing the final line without the use of a loop. I was thinking scatter_add_ or gather might help.

Just use torch.index_add function.
labels = torch.tensor([0, 1, 0])
x = torch.tensor([[1, 2, 3],[4, 5, 6],[7, 8, 9]])

nrow = torch.unique(labels).size(0)
ncol = x.size(1)
out = torch.zeros((nrow, ncol), dtype=x.dtype)
out.index_add_(0, labels, x)
print(out)

# the output will be 
# tensor([[ 8, 10, 12],
#         [ 4,  5,  6]])
 

