Custom Aggregation function

[rajajisai]

I am currently working on paper which has some convoluted computations. So i figured i would need a custom implementation of the aggregate function in MessagePassing module of PyTorch. I want to know if its possible to get the desired result in a using only vectorized computations instead of using data structures like lists ,dict etc.

Lets say the input to the aggregate function and the required output is as follows

So basically i don't want an aggregate or average or any of the reduction methods provided.I just want to append the feature values from source nodes appended against the destination index in resulting tensor.
This can be easily achieved using a list without padding .But using a list slowing speed for further computation. I need a resulting tensorwhich has the required messages from neighboring nodes appended and tensors of zeros padded wherever required. The tensor by itself should not take took long to compute
I think what i need is somewhat similar to mailbox in DGL. I would be glad if someone could help me here

[akihironitta]

Hi @rajajisai! I recently tried the same thing in my PR #7459. I hope this helps :)

import torch
from torch_geometric.utils import index_sort, to_dense_batch
src_feat = torch.tensor(
    [
        [1, 1, 1],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4],
        [5, 5, 5],
        [6, 6, 6],
    ]
)
dst_indices = torch.tensor([1, 2, 5, 5, 5, 1], dtype=torch.long)

# to_dense_batch assumes that the indices are sorted
_, indices = index_sort(dst_indices)
result, _ = to_dense_batch(src_feat[indices], dst_indices[indices])
print(result)

tensor([[[0, 0, 0],
         [0, 0, 0],
         [0, 0, 0]],

        [[1, 1, 1],
         [6, 6, 6],
         [0, 0, 0]],

        [[2, 2, 2],
         [0, 0, 0],
         [0, 0, 0]],

        [[0, 0, 0],
         [0, 0, 0],
         [0, 0, 0]],

        [[0, 0, 0],
         [0, 0, 0],
         [0, 0, 0]],

        [[3, 3, 3],
         [4, 4, 4],
         [5, 5, 5]]])

https://pytorch-geometric.readthedocs.io/en/2.3.1/modules/utils.html?highlight=to_dense_batch#torch_geometric.utils.to_dense_batch

[rajajisai]

@akihironitta you are a savior!! I am extremely thankful for the quick response. My only regret is not posting the query sooner :)