SageConv LSTM with Ordering

[barkincavdaroglu]

Hi,

While aggregating node's neighborhood, I want to order them using their edge weights. Since SageConv doesn't allow this, I created a custom class that handles this with for-loops as below:

seqs, lenghts = [], []
for node_idx in range(node_fts.shape[0]):
    seqs.append(
        dst_node_fts[
            torch.argsort(
                edge_fts[
                    :,
                ][edges[0] == node_idx],
                descending=True,
            )
        ]
    )
    lenghts.append(len(seqs[-1]))
lenghts = torch.tensor(lenghts)
seqs = pad_sequence(seqs, batch_first=True)
packed = nn.utils.rnn.pack_padded_sequence(
    seqs, lenghts.to("cpu"), batch_first=True, enforce_sorted=False
)

Is there a way to handle this by modifying SageConv? I thought about copy-pasting it and making changes there but I can't seem to find where SageConv aggregates neighborhoods to create sequences to pass into LSTM.

[wsad1]

For all conv operators aggregation is handled internally by an Aggregation object. You can create a new aggregation class and pass that as the aggr argument when initializing SAGEConv. Check out torch_geometric.nn.aggr for available aggregation operators.

[barkincavdaroglu]

Thank you so much!