Vectorised way to get graph-level outputs from node features without pooling?

[davidbuterez]

Hi,

First of all thanks for all the work put into this framework!

My issue is simple on the surface: I'd like to be able to use an arbitrary operator (e.g. MLP) to go from node features to graph-level outputs, as opposed to, for example, the global_<operator>_pool functions, which I see are efficiently implemented using scatter operations. For my purposes I'd like to be able to first group the node features back to their graph (using the batch indices vector). Using padding to a fixed size these graph-level outputs can then be used in a neural network. I have implemented this using loops and it seems to make a difference for the problem at hand. However, this is not efficient. Is there a vectorised alternative for this?

Thanks and apologies if there is an easy solution I missed!

[rusty1s]

We provide such functionality via torch_geometric.utils.to_dense_batch.