Retrieve corresponding node embedding after NeighborLoader sampler

[sidhantls]

I'm interested in using NeighborLoader to lower the complexity of a forward pass of the dataset in GNN and get node embeddings. However, wanted to know can I index the embedding of a node of concern after the sampler samples a batch? If I'm interested in a particular node embedding, I can retrieve its node embedding if i use the original dataset, by indexing the corresponding idx- but when i utilize the data sampler, how would i index the embedding of the same node?

Lets say I'm using the OGB_mag dataset. This is the data:

HeteroData(
  paper={
    x=[736389, 128],
    year=[736389],
    y=[736389],
    train_mask=[736389],
    val_mask=[736389],
    test_mask=[736389]
  },

I'm interested getting the node embeddings of paper for paper 10, 11. To do this, I'd just retrieve the 10th and 11th index in the paper output of model(data.x_dict, data.edge_index_dict). However, if I utilize the NeighborLoader, (and ensure the 10th and 11th index are part of the mask), I might get x of size: x=[20839, 128]. Now, how would I get the embedding of paper 10 and paper 11? Indexing 10 and 11 in the minibatch does not lead me to the 10th and 11th paper of interest originally.

Any suggestions on what can be done?

[wsad1]

One way to do this might be to add a variable that indicates original node id, as described here #3528, and then search for that node id in the neighbor loaders output. Although this might be slow.