Problem understanding max_pool_x layer

[Herijuana]

Hi,

I have problems understanding the "max_pool_x" layer in torch_geometric. I understand that it takes as arguments a cluster vector specifying which node corresponds to which cluster, the node-feature tensor x and the batch vector specifying which node corresponds to which instance of a batch. As output I am expecting the max along the node-dimension of the node-feature tensor over each cluster for each instance in the batch. But it seems to me the function is ignoring the batch. This makes sense, since the batching combines all graphs in a batch into a single graph and it seems to me that the function just takes the max along the node-dimension over each cluster of this graph. But I was expecting that the batch-vector is given as an argument to keep track of the instances in the batch.

If I would use global_max_pooling I would get a tensor of shape (batch_size,node_feature_size). If I use max_pool_x I get something of shape (number_clusters,node_feature_size).

So my question is, how can I take the batches into account such that the function does its thing for each instance in the batch separately ? I want to end up with a tensor of shape (batch_size*number_clusters,node_feature_size).

Thanks.

[rusty1s]

Yeah, max_pool_x pools both x and batch and returns a tuple, so it does not incorporate batch into the aggregation. You can easily achieve something like what you want via

cluster = ...
batch = ...
cluster = num_cluster * batch + cluster
global_max_pool(x, cluster)

[Herijuana]

Thank you for your answer. Could you also explain to me what the second output of the max_pool_x function does ? It takes the batch in and outputs a batch again but I don't understand whats happening.

[rusty1s]

max_pool_x takes in a cluster vector, and pools both node features x, and the batch vector batch based on it. As such, it returns a tuple of (pooled_x, pooled_batch).