Select¶

The most general interface is Select, which is the parent class for every \(\texttt{select}\) operator in tgp.

`Select`	An abstract base class implementing a sparse \(\texttt{select}\) operator that maps the nodes of an input graph to supernodes of the pooled one.
`SelectOutput`	The output of a `Select` method, which holds an assignment from selected nodes to their respective cluster(s).
`MLPSelect`	The \(\texttt{select}\) operator used by most of the dense pooling methods.
`DPSelect`	The Dirichlet Process selection operator for the `BNPool` operator, as proposed in the paper "BN-Pool: Bayesian Nonparametric Graph Pooling" (Castellana & Bianchi, 2025).
`EdgeContractionSelect`	The \(\texttt{select}\) operator from the papers "Towards Graph Pooling by Edge Contraction" (Diehl et al. 2019) and "Edge Contraction Pooling for Graph Neural Networks" (Diehl, 2019).
`EigenPoolSelect`	The \(\texttt{select}\) operator for EigenPooling.
`GraclusSelect`	The \(\texttt{select}\) operator inspired by the paper "Weighted Graph Cuts without Eigenvectors: A Multilevel Approach" (Dhillon et al., TPAMI 2007).
`IdentitySelect`	Identity select operator that maps each node to itself (no pooling).
`LaPoolSelect`	The select operator for the LaPool operator (`LaPooling`) as proposed in the paper Towards Interpretable Sparse Graph Representation Learning with Laplacian Pooling.
`KMISSelect`	Computes the node assignments following the Maximal \(k\)-Independent Set (\(k\)-MIS) algorithm, as defined in the paper "Generalizing Downsampling from Regular Data to Graphs" (Bacciu et al., AAAI 2023).
`MaxCutSelect`	The MaxCut \(\texttt{select}\) operator from the paper "MaxCutPool: differentiable feature-aware Maxcut for pooling in graph neural networks" (Abate & Bianchi, ICLR 2025).
`NDPSelect`	The select operator for Node Decimation Pooling (`NDPPooling`), as presented in the paper "Hierarchical Representation Learning in Graph Neural Networks with Node Decimation Pooling" (Bianchi et al., TNNLS 2020).
`NMFSelect`	Select operator that performs Non-negative Matrix Factorization pooling as proposed in the paper "A Non-Negative Factorization approach to node pooling in Graph Convolutional Neural Networks" (Bacciu and Di Sotto, AIIA 2019).
`TopkSelect`	The top-\(k\) \(\texttt{select}\) operator used by scoring-based pooling methods.
`SEPSelect`	Select operator for Structural Entropy Pooling (SEP).
`degree_scorer`
`eigenpool_select`	Compute EigenPool assignments and eigenvector pooling matrices.

class Select(*args, **kwargs)[source]¶

An abstract base class implementing a sparse \(\texttt{select}\) operator that maps the nodes of an input graph to supernodes of the pooled one.

It returns a SelectOutput containing the sparse supernode assignment matrix \(\mathbf{S} \in \mathbb{R}^{N \times K}\).

Forward pass.

Parameters:

x (Tensor, optional) – The node feature matrix of shape \([N, F]\), where \(N\) is the number of nodes in the batch and \(F\) is the number of node features. (default: None)
edge_index (Tensor, optional) – The edge indices. Is a tensor of of shape \([2, E]\), where \(E\) is the number of edges in the batch. (default: None)
edge_weight (Tensor, optional) – A vector of shape \([E]\) or \([E, 1]\) containing the weights of the edges. (default: None)
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. (default: None)
num_nodes (int, optional) – The total number of nodes of the graphs in the batch. (default: None)

Returns:

The output of \(\texttt{select}\) operator.

Return type:

SelectOutput

The output of a Select method, which holds an assignment from selected nodes to their respective cluster(s).

Parameters:

node_index (Tensor) – The indices of the selected nodes.
num_nodes (int) – The number of nodes.
cluster_index (Tensor) – The indices of the clusters each node in node_index is assigned to.
num_supernodes (int) – The number of clusters.
weight (Tensor, optional) – A weight vector, denoting the membership of a node to each cluster. (default: None)

property is_expressive: bool¶: Check if the assignment matrix is produced by an expressive pooling method. An assignment matrix is expressive if all rows sum to the same constant and that constant is non-zero.

property out_mask: Tensor | None¶

Boolean validity mask on pooled supernodes with shape \([B, K]\).

This is inferred from dense assignment matrix s and marks supernodes that have at least one assigned node. For s.dim() == 3 (\([B, N, K]\)), there is one mask row per graph. For s.dim() == 2 (\([N, K]\)) with batch, there is one row per graph id in batch. For s.dim() == 2 without batch, the result has shape \([1, K]\).

Returns None for sparse assignments.

property is_sparse: bool¶: Return True if s is a sparse tensor.

property is_dense: bool¶: Return True if s is a dense tensor.

property num_nodes: int¶: Return the number of input nodes represented by s.

property num_supernodes: int¶: Return the number of pooled nodes represented by s.

property node_index: Tensor | None¶: Return sparse row indices (node ids) when s is sparse.

property cluster_index: Tensor | None¶: Return sparse column indices (supernode ids) when s is sparse.

property weight: Tensor | None¶: Return sparse assignment values when s is sparse.

set_s_inv(method)[source]¶: Compute and store s_inv from s using the given strategy.

apply(func: Callable) → SelectOutput[source]¶: Applies func to tensors in s, s_inv, and tensor-valued extra attributes.

clone() → SelectOutput[source]¶: Performs a deep-copy of the object.

to(device: int | str, non_blocking: bool = False) → SelectOutput[source]¶: Performs tensor dtype and/or device conversion for both s and s_inv.

cpu() → SelectOutput[source]¶: Copies attributes to CPU memory for both s and s_inv.

cuda(device: int | str | None = None, non_blocking: bool = False) → SelectOutput[source]¶: Copies attributes to CUDA memory for both s and s_inv.

detach_() → SelectOutput[source]¶: Detaches attributes from the computation graph for both s and s_inv.

detach() → SelectOutput[source]¶: Detaches attributes from the computation graph by creating a new tensor for both s and s_inv.

requires_grad_(requires_grad: bool = True) → SelectOutput[source]¶: Tracks gradient computation for both s and s_inv.

assign_all_nodes(adj: Tensor | SparseTensor | None = None, weight: Tensor | None = None, max_iter: int = 5, batch: Tensor | None = None, closest_node_assignment: bool = True) → SelectOutput[source]¶

Extends a sparse selection to assign ALL nodes to the selected supernodes.

This method converts a sparse selection (where only a subset of nodes are initially selected, e.g. top-k selection) into a complete assignment where every node in the graph is assigned to one of the selected supernodes.

Parameters:

adj (Adj, optional) – Graph connectivity matrix. Can be an edge_index tensor of shape \((2, E)\) or SparseTensor. Required for "closest_node" strategy. (default: None)
weight (Tensor, optional) – Node-level weights for the assignment. Must have shape \((N,)\) where \(N\) is the total number of nodes. Note: This is different from edge weights. (default: None)
max_iter (int, optional) – Maximum number of message passing iterations for the "closest_node" strategy. Higher values allow assignment of more distant nodes through graph connectivity. Must be > 0 for "closest_node" strategy. (default: 5)
batch (Tensor, optional) – Batch assignment vector of shape \((N,)\) indicating which graph each node belongs to. When provided, ensures nodes are only assigned to supernodes within the same graph. (default: None)
closest_node_assignment (bool, optional) – If True, assign unlabeled nodes to the closest supernode. If False, use random assignment to supernodes. (default: True)

Returns:

A new SelectOutput with complete node-to-supernode assignments. The returned object has num_nodes assignments (one per node) and num_supernodes supernodes (same as the original selection).

Return type:

SelectOutput

Raises:

AssertionError – If adj is None for "closest_node" strategy.
AssertionError – If max_iter <= 0 for "closest_node" strategy.
ValueError – If weight size doesn’t match the number of nodes.
ValueError – If adj has an invalid type.
ValueError – If strategy is not recognized.

Example

>>> # Convert sparse top-k selection to full assignment
>>> # Assume we have a SelectOutput from top-k selection
>>> sparse_output = topk_selector(x, edge_index)  # Only k nodes selected
>>> print(sparse_output.node_index.size(0))  # k nodes
>>> # Extend to assign all nodes using graph connectivity
>>> full_output = sparse_output.assign_all_nodes(
...     adj=edge_index, closest_node_assignment=True, max_iter=5
... )
>>> print(full_output.node_index.size(0))  # N nodes (all nodes)
>>> print(full_output.num_supernodes)  # Still k supernodes

class MLPSelect(in_channels: int | List[int], k: int, batched_representation: bool = True, act: str | None = None, dropout: float = 0.0, s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The \(\texttt{select}\) operator used by most of the dense pooling methods.

It computes a dense assignment matrix \(\mathbf{S} \in \mathbb{R}^{B \times N \times K}\) from the node features \(\mathbf{X} \in \mathbb{R}^{B \times N \times F}\):

\[\mathbf{S} = \mathrm{softmax}(\texttt{MLP}(\mathbf{X}))\]

Parameters:

in_channels (int, list of int) – Number of hidden units for each hidden layer in the MLP used to compute cluster assignments. The first integer must match the size of the node features.
k (int) – Number of clusters or supernodes in the pooler graph.
batched_representation (bool, optional) – If True, expects batched input \(\mathbf{X} \in \mathbb{R}^{B \times N \times F}\) and returns assignment matrix \(\mathbf{S} \in \mathbb{R}^{B \times N \times K}\). If False, expects unbatched input \(\mathbf{X} \in \mathbb{R}^{N \times F}\) where \(N\) is the total number of nodes across all graphs, and returns assignment matrix \(\mathbf{S} \in \mathbb{R}^{N \times K}\). (default: True)
act (str or Callable, optional) – Activation function in the hidden layers of the MLP.
dropout (float, optional) – Dropout probability in the MLP. (default: 0.0)
s_inv_op (SinvType, optional) –
The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:
- "transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
- "inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

forward(x: Tensor, mask: Tensor | None = None, batch: Tensor | None = None, **kwargs) → SelectOutput[source]¶

Forward pass.

Parameters:

x (Tensor) – Node feature tensor. If batched_representation=True, expected shape is \(\mathbf{X} \in \mathbb{R}^{B \times N \times F}\), with batch-size \(B\), (maximum) number of nodes \(N\) for each graph, and feature dimension \(F\). If batched_representation=False, expected shape is \(\mathbf{X} \in \mathbb{R}^{N \times F}\), where \(N\) is the total number of nodes across all graphs in the batch.
mask (Tensor, optional) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. Only used when batched_representation=True. (default: None)
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. Only used when batched_representation=False. (default: None)

Returns:

The output of \(\texttt{select}\) operator.: If batched_representation=True, the assignment matrix \(\mathbf{S}\) has shape \(\mathbb{R}^{B \times N \times K}\). If batched_representation=False, the assignment matrix \(\mathbf{S}\) has shape \(\mathbb{R}^{N \times K}\).

Return type:

SelectOutput

class DPSelect(in_channels: int | List[int], k: int, batched_representation: bool = True, act: str | None = None, dropout: float = 0.0, s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The Dirichlet Process selection operator for the BNPool operator, as proposed in the paper “BN-Pool: Bayesian Nonparametric Graph Pooling” (Castellana & Bianchi, 2025).

DPSelect implements a Bayesian nonparametric selection mechanism to automatically learn both the number of clusters and their assignments through variational inference. The method uses a truncated stick-breaking representation of the Dirichlet Process to model cluster assignments:

\[v_{ik} \sim \text{Beta}(\alpha_{ik}, \beta_{ik}), \quad k = 1, \ldots, K-1, \quad i = 1, \ldots, N\]

where \(v_{ik}\) are the stick-breaking fractions. The assignment of node \(i\) to cluster \(k\) is computed as:

\[\pi_{ik} = v_{ik} \prod_{j=1}^{k-1} (1 - v_{ij}) \quad \text{for } k = 1, \ldots, K-1\]

The variational parameters \(\alpha_{ik}, \beta_{ik}\) are computed by an MLP from node features:

\[[\alpha_{i,1}, \ldots, \alpha_{i,K-1}, \beta_{i,1}, \ldots, \beta_{i,K-1}] = \text{softplus}(\text{MLP}(\mathbf{x}_i))\]

The procedure can be summarized as follows:

Feature Processing: Node features are processed by an MLP to produce \(2(K-1)\) outputs per node.
Parameter Extraction: The MLP output is split into \(\boldsymbol{\alpha}\) and \(\boldsymbol{\beta}\) parameters for the Beta distribution.
Sampling: Stick-breaking fractions are obtained from the sampling procedure implemented in Beta: \(v_{ik} = \text{Beta}(\alpha_{ik}, \beta_{ik}).rsample()\).
Cluster Assignment: The assignment matrix is computed via the stick-breaking construction.

Parameters:

in_channels (int, list of int) – Number of hidden units for each hidden layer in the MLP used to compute cluster assignments. The first integer must match the size of the node features.
k (int) – Maximum number of clusters \(K\). The actual number of active clusters is learned through the stick-breaking process.
batched_representation (bool, optional) – If True, expects batched input \(\mathbf{X} \in \mathbb{R}^{B \times N \times F}\) and returns assignment matrix \(\mathbf{S} \in \mathbb{R}^{B \times N \times K}\). If False, expects unbatched input \(\mathbf{X} \in \mathbb{R}^{N \times F}\) where \(N\) is the total number of nodes across all graphs, and returns assignment matrix \(\mathbf{S} \in \mathbb{R}^{N \times K}\). (default: True)
act (str or Callable, optional) – Activation function in the hidden layers of the MLP.
dropout (float, optional) – Dropout probability in the MLP. (default: 0.0)
s_inv_op (SinvType, optional) –
The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:
- "transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
- "inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

Note

This class extends MLPSelect but replaces the softmax assignment with the stick-breaking construction. The SelectOutput returned includes both the assignment matrix \(\mathbf{S}\) and the posterior distributions \(q(v_{ik})\) for computing KL divergence losses.

forward(x: Tensor, mask: Tensor | None = None, batch: Tensor | None = None, **kwargs) → SelectOutput[source]¶

Applies the Dirichlet Process selection operator to compute cluster assignments.

Parameters:

x (Tensor) – Node feature tensor. If batched_representation=True, expected shape is \(\mathbb{R}^{B \times N \times F}\). If batched_representation=False, expected shape is \(\mathbb{R}^{N \times F}\), where \(N\) is the total number of nodes across all graphs in the batch.
mask (Tensor, optional) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes. Only used when batched_representation=True. (default: None)
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. Only used when batched_representation=False. (default: None)

Returns:

The output of \(\texttt{select}\) operator.: If batched_representation=True, the assignment matrix \(\mathbf{S}\) has shape \(\mathbb{R}^{B \times N \times K}\). If batched_representation=False, the assignment matrix \(\mathbf{S}\) has shape \(\mathbb{R}^{N \times K}\).

Return type:

SelectOutput

class EdgeContractionSelect(in_channels: int, edge_score_method: Callable | None = None, dropout: float | None = 0.0, add_to_edge_score: float = 0.5, s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The \(\texttt{select}\) operator from the papers “Towards Graph Pooling by Edge Contraction” (Diehl et al. 2019) and “Edge Contraction Pooling for Graph Neural Networks” (Diehl, 2019). This implementation is based on the paper “Revisiting Edge Pooling in Graph Neural Networks” (Landolfi, 2022).

In short, a score is computed for each edge. Edges are contracted iteratively according to that score unless one of their nodes has already been part of a contracted edge.

Parameters:

in_channels (int) – Size of each input sample.
edge_score_method (callable, optional) – The function to apply to compute the edge score from raw edge scores. By default, this is the softmax over all incoming edges for each node. This function takes in a raw_edge_score tensor of shape [num_nodes], an edge_index tensor and the number of nodes num_nodes, and produces a new tensor of the same size as raw_edge_score describing normalized edge scores. Included functions are compute_edge_score_softmax(), compute_edge_score_tanh(), and compute_edge_score_sigmoid(). (default: compute_edge_score_softmax())
dropout (float, optional) – The probability with which to drop edge scores during training. (default: 0.0)
add_to_edge_score (float, optional) – A value to be added to each computed edge score. Adding this greatly helps with unpooling stability. (default: 0.5)

static compute_edge_score_softmax(raw_edge_score: Tensor, edge_index: Tensor, num_nodes: int) → Tensor[source]¶: Normalizes edge scores via softmax application.

static compute_edge_score_tanh(raw_edge_score: Tensor, edge_index: Tensor | None = None, num_nodes: int | None = None) → Tensor[source]¶: Normalizes edge scores via hyperbolic tangent application.

static compute_edge_score_sigmoid(raw_edge_score: Tensor, edge_index: Tensor | None = None, num_nodes: int | None = None) → Tensor[source]¶: Normalizes edge scores via sigmoid application.

forward(x: Tensor, edge_index: Tensor, **kwargs) → SelectOutput[source]¶

Parameters:

x (Tensor) – The node feature matrix of shape \([N, F]\), where \(N\) is the number of nodes in the batch and \(F\) is the number of node features.
edge_index (Tensor) – The edge indices. Is a tensor of of shape \([2, E]\), where \(E\) is the number of edges in the batch.
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs.

Returns:

The output of \(\texttt{select}\) operator.

Return type:

SelectOutput

class EigenPoolSelect(k: int, s_inv_op: Literal['transpose', 'inverse'] = 'transpose', num_modes: int = 5, normalized: bool = True)[source]¶

The \(\texttt{select}\) operator for EigenPooling.

This operator performs spectral clustering on the adjacency matrix to build a dense assignment matrix \(\mathbf{S} \in \{0,1\}^{N \times K}\) and the eigenvector pooling matrix \(\boldsymbol{\Theta} \in \mathbb{R}^{N \times (K\cdot H)}\) used by the EigenPooling reduce/lift steps.

The same assignment matrix may also be denoted as \(\boldsymbol{\Omega}\) in connectivity formulas; in this implementation \(\boldsymbol{\Omega} = \mathbf{S}\).

Parameters:

k (int) – Number of clusters (supernodes).
s_inv_op (SinvType, optional) – Operation used to compute \(\mathbf{S}_\text{inv}\) from \(\mathbf{S}\). (default: "transpose")
num_modes (int, optional) – Number of eigenvector modes \(H\). (default: 5)
normalized (bool, optional) – If True, use the normalized Laplacian. (default: True)

Forward pass.

Parameters:

x (Tensor, optional) – Node features (unused by EigenPooling). (default: None)
edge_index (Adj, optional) – Graph connectivity.
edge_weight (Tensor, optional) – Edge weights associated with edge_index. (default: None)
batch (Tensor, optional) – Batch vector for multi-graph inputs. (default: None)
num_nodes (int, optional) – Number of nodes in the graph. (default: None)

Returns:

Selection output with:

s: assignment matrix \(\mathbf{S}\)
theta: pooling matrix \(\boldsymbol{\Theta}\)

Return type:

SelectOutput

class GraclusSelect(s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The \(\texttt{select}\) operator inspired by the paper “Weighted Graph Cuts without Eigenvectors: A Multilevel Approach” (Dhillon et al., TPAMI 2007).

It implements a greedy clustering algorithm for picking an unmarked vertex and matching it with one of its unmarked neighbors (that maximizes its edge weight).

Parameters:

s_inv_op (SinvType, optional) –

The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:

"transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
"inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

forward(edge_index: Tensor, edge_weight: Tensor | None = None, num_nodes: int | None = None, **kwargs) → SelectOutput[source]¶

Forward pass.

Parameters:

edge_index (Tensor, optional) – The edge indices. Is a tensor of of shape \([2, E]\), where \(E\) is the number of edges in the batch. (default: None)
edge_weight (Tensor, optional) – A vector of shape \([E]\) or \([E, 1]\) containing the weights of the edges. (default: None)
num_nodes (int, optional) – The total number of nodes of the graphs in the batch. (default: None)

Returns:

The output of \(\texttt{select}\) operator.

Return type:

SelectOutput

class IdentitySelect[source]¶

Identity select operator that maps each node to itself (no pooling).

forward(*, edge_index: Tensor | SparseTensor | None = None, num_nodes: int | None = None, device: device | None = None, **kwargs) → SelectOutput[source]¶

Create identity mapping where each node maps to itself.

Parameters:

edge_index (Optional[Adj]) – The edge index of the graph.
num_nodes (Optional[int]) – The number of nodes in the graph. If not provided, it will be inferred from the edge_index.
device (Optional[torch.device]) – The device to use for the output tensors. If not provided, it will be inferred from the edge_index.

Returns:

The output of the identity select operator.

Return type:

SelectOutput

class LaPoolSelect(shortest_path_reg: bool = False, batched_representation: bool = True, s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The select operator for the LaPool operator (LaPooling) as proposed in the paper Towards Interpretable Sparse Graph Representation Learning with Laplacian Pooling. (Emmanuel Noutahi et al., 2019).

This operator computes a soft assignment matrix \(\mathbf{S}\) by first identifying a set of leaders, and then assigning every remaining node to the cluster of the closest leader:

\[\begin{split}\begin{align*} \mathbf{v} &= \| \mathbf{L} \mathbf{X} \|_d \\ \mathbf{i} &= \{ i \mid \mathbf{v}_i > \mathbf{v}_j, \forall j \in \mathcal{N}(i) \} \\ \mathbf{S}^\top &= \texttt{SparseSoftmax} \left( \beta \frac{\mathbf{X}\mathbf{X}_{\mathbf{i}}^\top}{\|\mathbf{X}\|\|\mathbf{X}_{\mathbf{i}}\|} \right) \end{align*}\end{split}\]

where:

\(\mathbf{L}\) is the Laplacian matrix of the graph,
\(\mathbf{X}\) is the input node feature matrix,
\(\beta\) is a regularization vector that is applied element-wise to the selection matrix.

Parameters:

shortest_path_reg (bool, optional) – If True, \(\beta\) is equal to the inverse of the shortest path between each node and its corresponding leader (this can be expensive since it runs on CPU). Otherwise \(\beta=1\).
batched_representation (bool, optional) – If True, expects batched input \(\mathbf{X} \in \mathbb{R}^{B \times N \times F}\) and a dense adjacency tensor \(\mathbf{A} \in \mathbb{R}^{B \times N \times N}\). If False, expects unbatched input \(\mathbf{X} \in \mathbb{R}^{N \times F}\) and a sparse adjacency in one of the formats supported by Adj (or a dense \([N, N]\) tensor). (default: True)
s_inv_op (SinvType, optional) –
The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:
- "transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
- "inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

Forward pass.

Parameters:

x (Tensor) – The node feature matrix of shape \([N, F]\), where \(N\) is the number of nodes in the batch and \(F\) is the number of node features.
edge_index (Adj, optional) – The connectivity matrix. It can either be a torch_sparse.SparseTensor of (sparse) shape \([N, N]\), where \(N\) is the number of nodes in the batch or a Tensor of shape \([2, E]\), where \(E\) is the number of edges in the batch. For batched dense inputs, it also accepts dense adjacency tensors of shape \([B, N, N]\) (or \([N, N]\) for unbatched dense inputs).
edge_weight (Tensor, optional) – A vector of shape \([E]\) or \([E, 1]\) containing the weights of the edges. (default: None)
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. (default: None)
mask (Tensor, optional) – Input-node validity mask \(\mathbf{M} \in {\{ 0, 1 \}}^{B \times N}\) with True on real (non-padded) nodes in each graph (batched mode only). (default: None)
num_nodes (int, optional) – The total number of nodes of the graphs in the batch. (default: None)

Returns:

The output of \(\texttt{select}\) operator.

Return type:

SelectOutput

class KMISSelect(in_channels: int | None = None, order_k: int = 1, scorer: str = 'linear', score_heuristic: str | None = 'greedy', force_undirected: bool = False, s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

Computes the node assignments following the Maximal \(k\)-Independent Set (\(k\)-MIS) algorithm, as defined in the paper “Generalizing Downsampling from Regular Data to Graphs” (Bacciu et al., AAAI 2023).

To compute the \(k\)-MIS, the algorithm greedily selects the nodes in their canonical order. If a permutation perm is provided, the nodes are extracted following that permutation instead.

Parameters:

in_channels (int, optional) – Size of each input sample. Ignored if scorer is not "linear". (default: None)
order_k (int) – The \(k\)-th order for the independent set. (default: 1)
scorer (str) –
A function that computes a score for each node. Nodes with higher score have a higher chance of being selected for pooling. It can be one of:
- "linear" (default): Uses a sigmoid-activated linear layer to compute the scores. in_channels must be set when using this option.
- "random": Assigns a random score in \([0, 1]\) to each node.
- "constant": Assigns a constant score of \(1\) to each node.
- "canonical": Assigns the score \(-i\) to the \(i\)-th node.
- "degree": Uses the degree of each node as the score.
score_heuristic (str, optional) –
Heuristic to increase the total score of selected nodes. Given an initial score vector \(\mathbf{s} \in \mathbb{R}^n\), options include:
- None: No heuristic applied.
- "greedy" (default): Computes the updated score \(\mathbf{s}'\) as
  
  \[\mathbf{s}' = \mathbf{s} \oslash (\mathbf{A} + \mathbf{I})^k \mathbf{1}\]
  
  where \(\oslash\) is element-wise division.
- "w-greedy": Computes the updated score \(\mathbf{s}'\) as
  
  \[\mathbf{s}' = \mathbf{s} \oslash (\mathbf{A} + \mathbf{I})^k \mathbf{s}\]
force_undirected (bool, optional) – Whether to force the input graph to be undirected. (default: False)
s_inv_op (SinvType, optional) –
The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:
- "transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
- "inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

Forward pass.

Parameters:

edge_index (Adj) – The connectivity matrix. It can either be a torch_sparse.SparseTensor of (sparse) shape \([N, N]\), where \(N\) is the number of nodes in the batch or a Tensor of shape \([2, E]\), where \(E\) is the number of edges in the batch.
edge_weight (Tensor, optional) – A vector of shape \([E]\) or \([E, 1]\) containing the weights of the edges. (default: None)
x (Tensor, optional) – The node feature matrix of shape \([N, F]\), where \(N\) is the number of nodes in the batch and \(F\) is the number of node features. (default: None)
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. (default: None)
num_nodes (int, optional) – The total number of nodes of the graphs in the batch. (default: None)

Returns:

The output of the \(\texttt{select}\) operator.

Return type:

SelectOutput

class MaxCutSelect(in_channels: int, ratio: int | float = 0.5, assign_all_nodes: bool = True, max_iter: int = 5, mp_units: list = [32, 32, 32, 32, 16, 16, 16, 16, 8, 8, 8, 8], mp_act: str = 'tanh', mlp_units: list = [16, 16], mlp_act: str = 'relu', act: str = 'tanh', delta: float = 2.0, min_score: float | None = None, s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The MaxCut \(\texttt{select}\) operator from the paper “MaxCutPool: differentiable feature-aware Maxcut for pooling in graph neural networks” (Abate & Bianchi, ICLR 2025).

This operator computes node scores using a specialized neural network that optimizes the MaxCut objective, then performs top-k selection based on these scores.

The MaxCut scoring process consists of:

Score Computation: A graph neural network computes node-level scores \(\mathbf{s} \in [-1, 1]^N\) via:

\[\mathbf{s} = \tanh(\text{MLP}(\text{GNN}(\mathbf{X}, \mathbf{A})))\]
Top-k Selection: Select top-k nodes based on scores:

\[\mathbf{i} = \text{top}_k(|\mathbf{s}|)\]

The computed scores are stored in the SelectOutput and can be accessed by the pooler for loss computation.

Parameters:

in_channels (int) – Size of each input feature.
ratio (Union[int, float]) – Graph pooling ratio for top-k selection. (default: 0.5)
assign_all_nodes (bool, optional) – Whether to create assignment matrices that map all nodes to the closest supernode (True) or perform standard top-k selection (False). (default: True)
max_iter (int, optional) – Maximum distance for the closest node assignment. (default: 5)
mp_units (list, optional) – List of hidden units for message passing layers. (default: [32, 32, 32, 32, 16, 16, 16, 16, 8, 8, 8, 8])
mp_act (str, optional) – Activation function for message passing layers. (default: "tanh")
mlp_units (list, optional) – List of hidden units for MLP layers. (default: [16, 16])
mlp_act (str, optional) – Activation function for MLP layers. (default: "relu")
act (str, optional) – Activation function for the final score. (default: "tanh")
delta (float, optional) – Delta parameter for propagation matrix computation. (default: 2.0)
min_score (float, optional) – Minimal node score threshold. Inherited from TopkSelect. (default: None)
s_inv_op (SinvType, optional) –
The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:
- "transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
- "inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

forward(x: Tensor, edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, **kwargs) → SelectOutput[source]¶

Forward pass of the MaxCut selector.

Parameters:

x (Tensor) – Node features of shape \((N, F)\).
edge_index (Tensor) – Graph connectivity in COO format of shape \((2, E)\).
edge_weight (Tensor, optional) – Edge weights of shape \((E,)\). (default: None)
batch (Tensor, optional) – Batch assignments of shape \((N,)\). (default: None)

Returns:

Selection output containing node indices, weights, and scores.

Return type:

SelectOutput

class NDPSelect(s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The select operator for Node Decimation Pooling (NDPPooling), as presented in the paper “Hierarchical Representation Learning in Graph Neural Networks with Node Decimation Pooling” (Bianchi et al., TNNLS 2020).

It partitions the nodes based on the sign of the largest eigenvector of the Laplacian. One side of the partition becomes the set of supernodes in the pooled graph, while the other side is dropped.

Parameters:

s_inv_op (SinvType, optional) –

The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:

"transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
"inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

forward(edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, *, batch: Tensor | None = None, num_nodes: int | None = None, **kwargs) → SelectOutput[source]¶

Forward pass.

Parameters:

edge_index (Adj, optional) – The connectivity matrix. It can either be a torch_sparse.SparseTensor of (sparse) shape \([N, N]\), where \(N\) is the number of nodes in the batch or a Tensor of shape \([2, E]\), where \(E\) is the number of edges in the batch.
edge_weight (Tensor, optional) – A vector of shape \([E]\) or \([E, 1]\) containing the weights of the edges. (default: None)
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. (default: None)
num_nodes (int, optional) – The total number of nodes of the graphs in the batch. (default: None)

Returns:

The output of \(\texttt{select}\) operator.

Return type:

SelectOutput

static eval_cut(total_volume, L, z)[source]¶

Computes the normalized size of a cut.

Parameters:

total_volume (float) – Total graph volume used to normalize the cut value.
L (scipy.sparse.csr_matrix) – The (unweighted) Laplacian.
z (ndarray) – Partition vector of shape \([N, 1]\) with entries in \(\{-1, 1\}\).

Returns:

A value in \([0,1]\) representing the normalized size of the cut.

Return type:

ndarray

static sign_partition(vec_or_size: Tensor | int) → Tuple[Tensor, Tensor][source]¶: Split indices into positive and negative partitions.

class NMFSelect(k: int, s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

Select operator that performs Non-negative Matrix Factorization pooling as proposed in the paper “A Non-Negative Factorization approach to node pooling in Graph Convolutional Neural Networks” (Bacciu and Di Sotto, AIIA 2019).

This select operator computes the non-negative matrix factorization

\[\mathbf{A} = \mathbf{W}\mathbf{H}\]

and returns \(\mathbf{H}^\top\) as the dense clustering matrix.

Parameters:

k (int) – Number of clusters or supernodes in the pooler graph.
s_inv_op (SinvType, optional) –
The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:
- "transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
- "inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

forward(edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, *, batch: Tensor | None = None, num_nodes: int | None = None, fixed_k: bool = False, **kwargs) → SelectOutput[source]¶

Forward pass of the select operator.

Parameters:

edge_index (Adj) – Graph connectivity. Sparse graph connectivity (edge_index, SparseTensor, or torch COO).
edge_weight (Tensor, optional) – Edge weights for sparse inputs. (default: None)
batch (Tensor, optional) – Batch vector for sparse inputs. (default: None)
num_nodes (int, optional) – Number of nodes for sparse inputs when it cannot be inferred from edge_index. (default: None)
fixed_k (bool, optional) – If True, force assignment width to exactly k for single sparse graphs by right-padding zero columns. Useful for pre-coarsening where per-sample outputs are collated together. (default: False)

Returns:

The output of \(\texttt{select}\) operator.

Return type:

SelectOutput

class TopkSelect(in_channels: int | None = None, ratio: int | float = 0.5, min_score: float | None = None, act: str | Callable = 'tanh', s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

The top-\(k\) \(\texttt{select}\) operator used by scoring-based pooling methods.

Behavior based on in_channels:

When in_channels is None or <= 1: The operator does not learn a projection vector and directly uses the input as node scores (optionally applying the activation function). This mode is useful when you want to provide pre-computed scores directly.
When in_channels is > 1: The operator learns a projection vector \(\mathbf{p}\) and computes scores by projecting the input features.

Score computation:

If min_score is None, computes:

\[\begin{split}\mathbf{s} &= \begin{cases} \sigma(\mathbf{x}) & \text{if } \texttt{in_channels} \leq 1 \\ \sigma \left( \frac{\mathbf{X}\mathbf{p}}{\| \mathbf{p} \|} \right) & \text{if } \texttt{in_channels} > 1 \end{cases}\end{split}\]

Then select the top-k nodes:

\[\mathbf{i} = \mathrm{top}_k(\mathbf{s})\]

If min_score is a value \(\tilde{\alpha} \in [0,1]\), computes:

\[\begin{split}\mathbf{s} &= \begin{cases} \mathrm{softmax}(\mathbf{x}, \text{batch}) & \text{if } \texttt{in_channels} \leq 1 \\ \mathrm{softmax}(\mathbf{X}\mathbf{p}, \text{batch}) & \text{if } \texttt{in_channels} > 1 \end{cases}\end{split}\]

Then select all nodes above the threshold:

\[\mathbf{i} = \{ j : \mathbf{s}_j > \tilde{\alpha} \}\]

where \(\mathbf{p}\) is the learnable projection vector and \(\sigma\) is the activation function.

Input handling:

2D input \(\mathbf{X} \in \mathbb{R}^{N \times F}\): Standard node feature matrix.
1D input \(\mathbf{x} \in \mathbb{R}^{N}\): When in_channels <= 1, used directly as scores. When in_channels > 1, reshaped to \(\mathbb{R}^{N \times 1}\) and projected.

Warning

When providing pre-computed scores, set act to "identity" or "linear" to avoid applying an activation function.

The SelectOutput contains a sparse assignment matrix \(\mathbf{S}\) that can be thought as dropping all the columns \(j \notin \mathbf{i}\) of the diagonal matrix \(\text{diag}(\mathbf{s})\).

Parameters:

in_channels (int, optional) – Size of each input sample. When None or <= 1, no learnable projection is used and input is treated as scores. When > 1, a learnable projection vector is used to compute scores from features. (default: None)
ratio (float or int) – The graph pooling ratio, which is used to compute \(k = \lceil \mathrm{ratio} \cdot N \rceil\), or the value of \(k\) itself, depending on whether the type of ratio is float or int. This value is ignored if min_score is not None. (default: 0.5)
min_score (float, optional) – Minimal node score \(\tilde{\alpha}\) which is used to compute indices of pooled nodes \(\mathbf{i} = \mathbf{s}_i > \tilde{\alpha}\). When this value is not None, the ratio argument is ignored. (default: None)
act (str or callable, optional) – The non-linearity \(\sigma\) to use when computing the score. Use "identity", "linear", or "none" to avoid applying any activation when providing pre-computed scores. (default: "tanh")
s_inv_op (SinvType, optional) –
The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix \(\mathbf{S}\). \(\mathbf{S}_\text{inv}\) is stored in the "s_inv" attribute of the SelectOutput. It can be one of:
- "transpose" (default): Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).
- "inverse": Computes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^+\), the Moore-Penrose pseudoinverse of \(\mathbf{S}\).

Examples

Using with node features (learnable projection):

>>> import torch
>>> from tgp.select.topk_select import TopkSelect
>>> # Node feature matrix: 5 nodes, 3 features each
>>> x = torch.randn(5, 3)
>>> selector = TopkSelect(in_channels=3, ratio=0.6)
>>> output = selector(x)
>>> print(f"Selected {output.num_supernodes} out of {output.num_nodes} nodes")

Using with pre-computed scores:

>>> # Pre-computed node scores
>>> scores = torch.tensor([0.1, 0.8, 0.3, 0.9, 0.2])
>>> selector = TopkSelect(in_channels=None, ratio=0.4, act="identity")
>>> output = selector(scores)
>>> print(f"Top nodes: {output.node_index}")  # Should select nodes 1 and 3

forward(x: Tensor, *, batch: Tensor | None = None, **kwargs) → SelectOutput[source]¶

Forward pass.

Parameters:

x (Tensor) – The node feature matrix of shape \([N, F]\) or node scores of shape \([N]\), where \(N\) is the number of nodes in the batch and \(F\) is the number of node features.
batch (Tensor, optional) – The batch vector \(\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N\), which indicates to which graph in the batch each node belongs. (default: None)

Returns:

The output of the \(\texttt{select}\) operator.

Return type:

SelectOutput

class SEPSelect(s_inv_op: Literal['transpose', 'inverse'] = 'transpose')[source]¶

Select operator for Structural Entropy Pooling (SEP).

The selector builds a coding tree per graph and uses its depth-1 partitions as hard cluster assignments.

Parameters:: s_inv_op (SinvType, optional) – The operation used to compute \(\mathbf{S}_\text{inv}\) from the select matrix. (default: "transpose")

Forward pass.

Parameters:

x (Tensor, optional) – Unused placeholder to keep interface compatibility.
edge_index (Adj) – Graph connectivity.
edge_weight (Tensor, optional) – Edge weights of shape \([E]\) or \([E, 1]\).
batch (Tensor, optional) – Batch vector assigning nodes to graphs.
num_nodes (int, optional) – Total number of nodes in the input batch.

Returns:

Hard assignment from nodes to depth-1 SEP clusters.

Return type:

SelectOutput

multi_level_select(edge_index: Tensor | SparseTensor | None = None, edge_weight: Tensor | None = None, *, batch: Tensor | None = None, num_nodes: int | None = None, levels: int = 1, **kwargs) → list[SelectOutput][source]¶: Compute multiple sequential SEP selections from a single tree build.

degree_scorer(edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, num_nodes: int | None = None, dim: int = 1)[source]¶

eigenpool_select(edge_index: Tensor | SparseTensor, k: int, edge_weight: Tensor | None = None, batch: Tensor | None = None, num_nodes: int | None = None, fixed_k: bool = False, s_inv_op: Literal['transpose', 'inverse'] = 'transpose', num_modes: int = 5, normalized: bool = True) → SelectOutput[source]¶

Compute EigenPool assignments and eigenvector pooling matrices.

Given a graph with \(N\) nodes, this function computes:

a hard assignment matrix \(\mathbf{S} \in \{0,1\}^{N \times K}\) via spectral clustering;
an eigenvector pooling matrix \(\boldsymbol{\Theta} \in \mathbb{R}^{N \times (K\cdot H)}\).

For consistency with the connector notation, \(\boldsymbol{\Omega}\) used in EigenPoolConnect is the same matrix as \(\mathbf{S}\).

Parameters:

edge_index (Adj) – Graph connectivity in edge index or dense adjacency format.
k (int) – Number of clusters (supernodes).
edge_weight (Tensor, optional) – Edge weights associated with edge_index. (default: None)
batch (Tensor, optional) – Batch vector \(\mathbf{b} \in \{0,\dots,B-1\}^N\) for multi-graph inputs. (default: None)
num_nodes (int, optional) – Total number of nodes. Useful when edge_index is empty. (default: None)
fixed_k (bool, optional) – If True, always use exactly k output clusters (allowing empty clusters). If False, single-graph mode may reduce the effective number of clusters for tiny graphs. (default: False)
s_inv_op (SinvType, optional) – Operation used to compute \(\mathbf{S}_\text{inv}\) stored in SelectOutput. (default: "transpose")
num_modes (int, optional) – Number of eigenvector modes \(H\) used to build \(\boldsymbol{\Theta}\). (default: 5)
normalized (bool, optional) – If True, uses the normalized Laplacian for eigenvectors. (default: True)

Returns:

Selection output with:

s: dense one-hot assignment matrix \([N, K]\)
theta: pooling matrix \(\boldsymbol{\Theta}\) (or a list of per-graph matrices for multi-graph batches)

Return type:

SelectOutput