Pytorch: all-but-one summation?
I am in the process of moving some message passing code from Numpy to Pytorch. I am unsure of how to do this single step of a much larger algorithm. Below is the simplest explanation of the step.
Given the following:
index = [[2,0,1], [2,2,0]]
value = [[0.1, 1.2, 2.3], [3.4, 4.5, 5.6]]

I would like to compute the &quot;all-but-one&quot; sum of messages to each index.  Here is a graphical representation:

The answer I am looking for is:
ans = [[7.9, 5.6, 0], [4.6, 3.5, 1.2]]

The explanation is that, for example, index[0][0] points at node 2. The sum of all messages at node 2 is 0.1+3.4+4.5=8. However we want to exclude the message we are considering (value[0][0]=0.1), so we obtain ans[0][0]=7.9. If only a single index points at a node then the answer is 0 (eg. node 1 with ans[0][2]).
I would be happy with computing the sums for each node, and then subtracting out the individual messages.  I am aware that this can lead to loss of significance, but I believe that my use case is very well-behaved (eg. no floating point infinities).
I can also provide the minimal numpy code, but even the minimal example is a bit long. I have looked at pytorch's scatter and gather commands, but I don't think that they're appropriate here.

After working with pytorch for a while longer, and writing some code for other situations, I realized that there is a much more efficient solution that I hadn't considered. So I am pasting it here for anyone else who comes after me:
import torch

index = [[2, 0, 1], [2, 2, 0]]
value = [[0.1, 1.2, 2.3], [3.4, 4.5, 5.6]]

# convert to tensor
index_tensor = torch.tensor(index)
value_tensor = torch.tensor(value)

num_nodes = 3
totals = torch.zeros(num_nodes)
totals = totals.index_add_(0, index_tensor.flatten(), value_tensor.flatten())
result = totals[index_tensor] - value_tensor
print(result)


It uses much less memory than the scatter_add solution given by Jodag. It avoids all the for loops given in the other solutions. Hooray for much faster code!
