Uncertain About Inputs for nn.avg_pool_x

[kevin3567]

For the torch_geometric.nn.avg_pool_x(cluster, x, batch, size=None) function, I understand that the function will return a tensor of shape (cluster_count, x_hidden_dim), as the function averages the rows in x which are from the same clusters.

However, what does the "batch" argument do? It does no seem to contribute to the function, but is a mandatory argument. Additionally, an explanation of the "size" argument would also be appreciated.

Thanks in advance.

[rusty1s]

Besides pooling of node features x, we also pool the batch vector to only contain the nodes of the coarsened graph.

The size argument can be used in case each example across a batch has the same number of clusters, and you want to reshape the output after pooling to [batch_size, num_clusters, num_features]. In that case, you specify size=num_clusters.