Sample one node from each graph in Batch object, given the node probabilities

[moe-assal]

I have a GNN that supports batching. Given (node_features, edge_index, batch) it outputs node_probabilites. How can I sample from those node_probabilites in a non-sequential manner?
I thought of using the ptr attribute (sequentially), but I'm sure there's a more efficient solution. Issue #1161 addresses this problem with no meaningful solution. It's been 4 years now though.
I am working with Batch objects.

[moe-assal]

I'm not sure if this is the "best" solution, but here:
Use Gumbel re-parameterization trick:

• transform probabilities into log probabilities (log p)
• sample a vector from the Gumbel Distribution (v)
• add them (log p + v)
• use scatter with a reduction='max' to get the maximum of each graph (argmax log p + v)
• Now you have the indexes of the sampled nodes

I think it's a good idea to incorporate something like that into the Batch class