Why is sorting edge_index necessary for PyG's LSTM Aggregation?

[Panichito]

Hi there, I've been doing Node embedding using GraphSAGE. Using the PyG library for Graph Neural Network, I decided to use the LSTM function for aggregate. The arguments it requires are as follows:

Input 1: Node features (|V|, F_in) - Here I use Node coordinates x-y in 2D plane (V x 2) and already normalized into the range of [0, 1] e.g.

          x         y
0  0.374540  0.598658
1  0.950714  0.156019
2  0.731994  0.155995

Input 2: Edge indices (2, |E|) - Adjacency matrix (V x V) but retrieves only the edge into (2, |E|) from the original adjacency matrix I have e.g.

idx   0  1  2
0   [[0, 1, 1], 
1    [1, 0, 1], 
2    [1, 1, 0]]

Above we have a shape (V x V) with 6 edges in the graph. We had to transform it a bit to accommodate PyG's use of shape (2, |E|) and I'd like to call it edge_index where edges are (0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1).

idx   0  1  2  3  4  5
0   [[0, 0, 1, 1, 2, 2],
1    [1, 2, 0, 2, 0, 1]]

Output: Node features (|V|, F_out) - Similar to Node coordinates, but they are not in 2D anymore, they are in a new embedding dimension with F_out dimensions.

My problem is that when using the LSTM aggregator it is forced to sort edge_index (Edge indices in input2) otherwise it will show an error saying ValueError: Can not perform aggregation since the 'index' tensor. is not sorted.

So I have to do sorting gives it with the following command:

# inside def __init__()
self.GraphSageLayer = SAGEConv(in_channels=2, out_channels=hidden_dim, aggr='lstm')

# inside def forward()
sorted_edge_index, _ = torch.sort(edge_index, dim=1)  # for LSTM
x = self.GraphSageLayer(coord.view(-1, 2), sorted_edge_index)  # using GraphSAGE

The sorted_edge_index tensor will look like this after sorting.

idx   0  1  2  3  4  5
0   [[0, 0, 1, 1, 2, 2],
1    [0, 0, 1, 1, 2, 2]]

I noticed that in my full-mesh graph of 3 nodes connected when I sorted it, the edges could be reinterpreted as (0, 0), (0, 0), (1, 1), (1, 1), (2, 2), (2, 2) which made me curious. And my questions are the following 2 things.

1. Why is sorting the edge_index necessary for LSTM? Will there be a problem with this?
2. How will my model determine which nodes are related after I sort edge_index in this manner? thus there are no longer any of the initial edge connection pairings. It is analogous to giving input edge pairs that are not present in the graph.

Ref. https://stackoverflow.com/questions/77987416/why-does-lstm-aggregation-in-pyg-need-to-sort-edge-index

[wsad1]

Why is sorting the edge_index necessary for LSTM? Will there be a problem with this?

Since graphs are unordered we need to sort the neighbors of a node to pass it to an LSTM.

How will my model determine which nodes are related after I sort edge_index in this manner? thus there are no longer any of the initial edge connection pairings. It is analogous to giving input edge pairs that are not present in the graph.

Use sort_edge_index or Data.sort() to correctly sort edge indicies as described in the LSTM aggregation document

[Panichito]

Got it, thank you!

    def fixed_index_of_edge_index(self, tensor):
        """
        Adjusts the indices of a given tensor based on the number of nodes.

        Args:
            tensor (list): The input tensor.

        Returns:
            list: The tensor with adjusted indices.
        """
        for i in range(len(tensor[0])):
            step = tensor[0][i] // self.num_nodes
            tensor[1][i] += step * self.num_nodes
        return tensor

    def forward(self, adj, dist, coord):
        """
        Args:
            adj (torch.Tensor): Adjacency matrix of shape (B x V x V), where B is the batch size and V is the number of nodes.
            dist (torch.Tensor): Distance matrix of shape (B x V x V), where B is the batch size and V is the number of nodes.
            coord (torch.Tensor): Coordinate matrix of shape (B x V x 2)
        """
        # Node embedding
        edge_index = adj.reshape(-1, self.num_nodes).eq(1).nonzero().t().contiguous()  # B x V x V -> (B*V) x V -> 2 x |E|
        edge_index = self.fixed_index_of_edge_index(edge_index)

        sort_edge_index = sort_edge_index(edge_index, num_nodes=self.num_nodes, sort_by_row=False)
        x = self.graphsage(coord.view(-1, 2), sort_edge_index)

[rusty1s]

It wouldn't be strictly necessary to required sorted indices. Our current implementation expects it so in order to accelerate the aggregation.