torch_geometric.nn.pool.max_pool_x

max_pool_x(cluster: Tensor, x: Tensor, batch: Tensor, batch_size: Optional[int] = None, size: Optional[int] = None) → Tuple[Tensor, Optional[Tensor]][source]

Max-Pools node features according to the clustering defined in cluster.

Parameters:

  • cluster (torch.Tensor) – The cluster vector \(\mathbf{c} \in \{ 0, \ldots, N - 1 \}^N\), which assigns each node to a specific cluster.

  • x (Tensor) – The node feature matrix.

  • batch (torch.Tensor) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which assigns each node to a specific example.

  • batch_size (int, optional) – The number of examples \(B\). Automatically calculated if not given. (default: None)

  • size (int, optional) – The maximum number of clusters in a single example. This property is useful to obtain a batch-wise dense representation, e.g. for applying FC layers, but should only be used if the size of the maximum number of clusters per example is known in advance. (default: None)

Return type:

(torch.Tensor, torch.Tensor) if size is None, else torch.Tensor