Pytorch differentiable conditional (index-based) sum
I have an idx array like [0, 1, 0, 2, 3, 1] and another 2d array data like the following:

[[0,  1,  2],
 [3,  4,  5],
 [6,  7,  8],
 [9,  10, 11],
 [12, 13, 14],
 [15, 16, 17]]


I want my output to be 4x3 in which 4 is the max of idx and 3 is the feature size (data.shape[1]) and in the output each element is the sum of features with corresponding index in the idx. Then output in this example would be:

[[6,  8,  10],
 [18, 20, 22],
 [9,  10, 11],
 [12, 13, 14]]


I can do it with iterating over range(3) and creating a mask on data and summing up them but it's not differentiable (I suppose). Is there any function(s) in Pytorch for this purpose? something like scatter().

Update: It seems I am looking for something named scatter sum which is implemented in this repository.

You are looking for index_add_:

import torch

x = torch.tensor([[ 0.,  1.,  2.],
        [ 3.,  4.,  5.],
        [ 6.,  7.,  8.],
        [ 9., 10., 11.],
        [12., 13., 14.],
        [15., 16., 17.]], dtype=torch.float)
idx = torch.tensor([0, 1, 0, 2, 3, 1], dtype=torch.long)  # note the dtype here, must be "long"
# init the sums to zero
y = torch.zeros((idx.max()+1, x.shape[1]), dtype=x.dtype)

# do the magic
y.index_add_(0, idx, x)


Gives the desired output


tensor([[ 6.,  8., 10.],
        [18., 20., 22.],
        [ 9., 10., 11.],
        [12., 13., 14.]])


