Use adjacency matrix to sum tensors' values with scatter_add

Hello. Is there a simple way to sum 2D tensors’ values with torch.scatter_add() by starting from the adjacency matrix? This is an example explaining what I would like to obtain:

x = torch.tensor([[1., 2.],
                  [3., 4.],
                  [5., 6.],
                  [7., 8.]])

adj = torch.tensor([[1, 1, 0, 1],
                    [1, 1, 1, 0],
                    [0, 1, 1, 0],
                    [1, 0, 0, 1]])

The output tensor after the operations should be

[[11., 14.],
 [9., 12.],
 [8., 10.], 
 [8., 10.]]

Meaning that sums should be computed in this way:

[ [1+3+7, 2+4+8],
  [3+1+5, 6+2+4],
  [5+3, 6+4],
  [7+1, 8+2]]

I remark that solving this with scatter_add() is quite important (but, of course, I am open to other solutions!).

It’s not scatter_add, but your task is essentially matrix multiplication adj.float() @ x.

Best regards

Thomas

1 Like

Ops, right, I didn’t see it at the beginning. Thanks!