[FIXED] How GNN layers can know which graph in a batch has the given node & node feature?

Issue

When we pass input as node features (x) and edge index (edge_index) to pytorch_geometric layer (e.g. GATConv), I am worried whether the layer can differentiate which batch sample the given node elements belong to.

x follows the shape [num of nodes, feature size] and edge_index follows shape [2, num of edges]. However, these 2 do not have the given information to know which input graph of batch size 32 have given node feature in the x.

Anyone can clarify on this ?

Solution

PyTorch-Geometric treats all the graphs in a batch as a single huge graph, with the individual graphs disconnected from each other. The node indices correspond to nodes in this big graph. This means there is no need for a batch dimension in x or edge_index.

Answered By - user58338