Message passing using sparse tensors

Hi Everyone,

I am trying to implement a Message Passing NN without pytorch-geometric. I find it quite easy to do using normal tensors:
torch.version

a = [[1., 2.], [3., 4.]]
b = [[5., 6.], [7., 8.], [9., 10.]]
ta = torch.tensor(a)
tb = torch.tensor(b).t()
torch.matmul(ta, tb)

> tensor([[17., 23., 29.],
>        [39., 53., 67.]])

tba = torch.mul(ta.view((2, 2, 1), 3, 0), tb.view((1, 2, 3), 3, 0))
tba.shape

> torch.Size([2, 2, 3])

tba.sum(1)

> tensor([[17., 23., 29.],
>        [39., 53., 67.]])

This way, I can replace sum with any other symmetric function (say max)

However, there is no way to handle a real graph’s adjaxcency matrix in a dense tensor. I cannot find a way to do the same with sparse tensors. Help?