Hi,
PyTorch provides the scatter_add operation, which can be very useful, but unfortunately is not deterministic. A deterministic mode for this operation is AFAIK not yet implemented (see this issue 3).
While the built-in scatter_add is very general, I’m wondering if there are solutions for the special case  where the index tensor is sorted. This could be a very relevant usecase where the index is just a segmentation/mask for variable-sized batches that were stacked together.
I’m thinking of something like this:
>>> batch_tensor = torch.tensor([1, 2, 3, 4, 5])
>>> index_tensor = torch.tensor([0, 0, 0, 1, 1])

# either
>>> scatter_add(batch_tensor, index_tensor)
tensor([6, 9])

# or
>>> torch.zeros(2).scatter_add(batch_tensor, index_tensor)
tensor([6, 9])

we have a batch with 2 elements, [1, 2, 3] and [4, 5], and we’re doing a sum global pool.
It seems there already are thoughts about sorted indices out there (e.g. point 3 here 1), but I don’t have enough insight to assess the current status. I’m wondering:

Is there already something like this in torch?
Is there already something like this in cuda?
If not, how straightforward is it to implement this?

Sure one can implement this directly in python by iterating over the index array, but I guess this would be painfully slow and there must be a better way that leverages the GPU.
I think a solution to this would have an immediate benefit for graph neural networks, but potentially also for other areas like CNNs with variable-sized images.
Stan.