Pytorch: all-but-one summation?

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!

Here is an approach which only requires a loop to iterate over the number of nodes, i.e. the maximum value in index plus 1 (in this case 3). It is not clear by the question if this value is always at most the number of columns in your input tensors so we just compute it explicitly.

This approach also allows you to compute gradients w.r.t. value_tensor and should be easily translatable to numpy as well.

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)

# optionally require gradients for value_tensor
# value_tensor.requires_grad_(True)

# perhaps this is always index_tensor.shape[1]? not clear from question
num_nodes = index_tensor.max() + 1

# compute total sum for each node
total_sum = torch.empty(num_nodes, device=value_tensor.device)
for n in range(num_nodes):
    total_sum[n] = value_tensor[index_tensor == n].sum()

# compute all-but-one
result = total_sum[index_tensor] - value_tensor

print(result)

which results in

tensor([[7.9000, 5.6000, 0.0000],
        [4.6000, 3.5000, 1.2000]])

Alternative using Tensor.scatter_add_

Here's an interesting alternative to the above which avoids the loop entirely by using the built in scatter-add operation. It requires more memory and is probably more difficult to follow than the previous solution but will likely be faster in some cases.

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)

# optionally require gradients for value_tensor
# value_tensor.requires_grad_(True)

num_rows = index_tensor.shape[0]
# perhaps this is always index_tensor.shape[1]? but doesn't need to be
num_nodes = index_tensor.max() + 1

# scatter_add will be applied to this tensor of zeros
scattered = torch.zeros((num_rows, num_nodes), device=value_tensor.device)

# apply scatter_add_
scattered.scatter_add_(1, index_tensor, value_tensor)

# which is equivalent to:
# num_cols = index_tensor.shape[1]
# for r in range(num_rows):
#     for c in range(num_cols):
#         scattered[r][index_tensor[r, c]] += value_tensor[r, c]

# sum the rows to get the total sum
total_sum = scattered.sum(dim=0)

# compute all-but-one
result = total_sum[index_tensor] - value_tensor

print(result)

Not sure if this is an improvement over your current solution, but you can do something like this:

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)

# initialize a tensor to store the result
ans = torch.empty_like(value_tensor)

# sum
for i, v_row in enumerate(value):
    for j, v in enumerate(v_row):
        ans[i, j] = value_tensor[index_tensor == index_tensor[i, j]].sum() - v

print(ans)
# tensor([[7.9000, 5.6000, 0.0000],
#         [4.6000, 3.5000, 1.2000]])

# if you need a list, just use ans.tolist()

Pretty sure there is a way to remove at least one of these for loops. I'll update the answer if I can figure it out.