I can’t come up with a better solution that this one:
adj_nz = adj.nonzero()
torch.stack([emb[adj_nz[adj_nz[:, 1] == u, 0]].sum(0) for u in torch.unique(adj_nz, sorted=True)])

Since the code might be a bit complicated, here are some explanations:


adj_nz is gives a tensor containing indices of all non zero values for each dim

tensor([[0, 0],
        [0, 1],
        [1, 2],
        [2, 0],
        [3, 0],
        [3, 2]])



adj_nz[:, 1] == u gives all rows for the current unique value in adj_nz. For u=0:

tensor([1, 0, 0, 1, 1, 0], dtype=torch.uint8)


We are using this to index adj_nz again and get the corresponding non-zero row indices using adj_nz[adj_nz[:, 1] == u, 0]:

tensor([0, 2, 3])


Once this is done, we can use this result to index emb and sum in dim0 using emb[adj_nz[adj_nz[:, 1] == u, 0]].sum(0):

tensor([18, 21, 24])


This has to be repeated for every unique value in adj_nz, which involves a list comprehension and a call to torch.stack. I’m not very happy about this last fact, but am currently unsure how to avoid it or if it might be a bottleneck at all.