PyTorch vectorise sum lookup quantity into buckets
Using PyTorch, I have figured out the following code for calculating totals of an item's property by some "bucket index":

DATASET_SIZE = 10
NUM_BUCKETS = 4
bucket_assignment = torch.tensor([0,1,2,3,0,1,2,3,0,1], dtype = torch.long)
values_to_add = torch.tensor([1,2,3,4,5,6,7,8,9,10], dtype = torch.float)
buckets = torch.zeros(NUM_BUCKETS, dtype = torch.float)
buckets.index_add_(0, bucket_assignment, values_to_add)

# Buckets is now tensor([15., 18., 10., 12.])


In my case this is specifically to check allocation bounds on a problem, and later code checks that no bucket is under- or over- allocated.

I would like to check multiple different possible assignments at once (and later pick a best option, code not shown). I thought I could do this by adding another dimension to bucket_assignment plus to buckets and have each row be a different set of assignments. However, this does not work as intended, because the second argument of index_add_ must be a simple vector, I cannot pass in any higher rank tensor.

E.g.

BATCH_SIZE = 2
DATASET_SIZE = 5
NUM_BUCKETS = 3
bucket_assignment = torch.tensor([[0,1,2,0,1], [1,1,1,2,1]], dtype = torch.long)
values_to_add = torch.tensor([1,2,3,4,5], dtype = torch.float)
buckets = torch.zeros(BATCH_SIZE, NUM_BUCKETS, dtype = torch.float)
buckets.index_add_(0, bucket_assignment, values_to_add)


I would like to get this result:

tensor([[5., 7., 3.], [ 0., 11.,  4.]])


Instead, I get an error:

RuntimeError: invalid argument 3: Index is supposed to be a vector at ../aten/src/TH/generic/THTensorEvenMoreMath.cpp:733


That's not unexpected due to the limitations of .index_add, but I don't know how to progress.

I am not sure what other approach would allow me to solve this problem in PyTorch - is there some other torch method I could use that would allow me to achieve the same thing. The main goal here is vectorisation and avoiding loops in Python, as in reality the batch sizes are large and I will be taking advantage of GPU acceleration.

If the batch size is the problem you could use torch.masked_select to get the values to add up for each bucket torch.masked_select(values_to_add, bucket_assignment == bucket_num), where PyTorch will broadcast the values_to_add and then only iterate over the buckets in plain python like so:

def bucket_sizes(bucket_num):
    mask = bucket_assignment == bucket_num
    buckets = torch.masked_select(values_to_add, mask)
    buckets = torch.split(buckets, list(mask.sum(dim=1)))
    return [bucket.sum() for bucket in buckets]

torch.tensor([bucket_sizes(i) for i in range(NUM_BUCKETS)]).T

