torch_geometric expand?

[ldv1]

Hi,

After a global pooling operation, I would like to "expand" the result to the graphs in the batch. Stated otherwise, I need the opposite of scatter. To illustrate, something like torch_geometric.XYZ.expand in

def forward(self, x, batch):
    global_max = global_max_pool(x, batch=batch)        
    y = torch_geometric.XYZ.expand(global_max, batch=batch)
    return x*y

In this case, the global max vector for a given graph is simply repeated for each node of the graph.

Is such a function already available?
Thanks for any help.

[rusty1s]

You can just do

global_max = global_max_pool(x, batch=batch)        
y = global_max[batch]

[ldv1]

Many thanks!