Unable to pool separately over subgraphs sampled by NeighborLoader

[egressybeni]

When using the NeighborLoader, I would like to be able to do graph pooling over the individual sampled subgraphs. This would be possible if batch was implemented for NeighborLoader. Currently batch is unavailable so when I do graph pooling I get a single value instead of batch_size values as an output.
This functionality would be useful whenever we want to make a prediction for the center node, but based on the final embeddings of the whole k-hop neighborhood. We might for example mark the center node and want to make a prediciton for the whole graph.
Similarly this would be useful for LinkNeighborLoader, where we want to make an edge prediction based on the final embeddings of the whole k-hop neighborhood.

Am missing an "obvious" way to achieve this?
Otherwise I might turn this into a feature request.

data_loader = NeighborLoader(train_data, batch_size=64)
data = next(iter(data_loader))
x = global_add_pool(data.x, data.batch) # unfortunately data.batch=None so this gives an output with shape[0] = 1 instead of shape[0] = batch_size = 64

(Of course I could use batch_size = 1, but that's not a great solution in general)

[rusty1s]

We are currently working on integrating this feature via the addition of pyg-lib. All we need to do on PyG side is to expose the disjoint argument, see here. Happy to take a PR :)

[egressybeni]

Thanks! I made a PR :) #5775