SRCPooling¶

The most general interface is SRCPooling, which is the parent class for every pooling operator in tgp.

`SRCPooling`	A base class for pooling layers based on the paper "Understanding Pooling in Graph Neural Networks" (Grattarola et al., TNNLS 2022).
`DenseSRCPooling`	A base class for dense pooling layers that extends `SRCPooling`.
`PoolingOutput`	The pooling output of a model of class `SRCPooling`.

class SRCPooling(selector: Select | None = None, reducer: Reduce | None = None, lifter: Lift | None = None, connector: Connect | None = None, cached: bool = False)[source]¶

A base class for pooling layers based on the paper “Understanding Pooling in Graph Neural Networks” (Grattarola et al., TNNLS 2022). Each pooler should inherit from this class.

SRCPooling decomposes a pooling layer into three components:

Select defines how input nodes map to supernodes.
Reduce defines how input node features are aggregated.
Lift defines how pooled node features are un-pooled.
Connect decides how the supernodes are connected to each other.

This class should return an object of type PoolingOutput.

Parameters:

selector (Select) – The node selection operator.
reducer (Reduce) – The node feature aggregation operator.
lifter (Lift) – The node feature un-pooling operator.
connector (Connect) – The edge connection operator.
cached (bool, optional) – If set to True, will cache the Select output and the Connect output. (default: False)

select(**kwargs) → SelectOutput[source]¶

Calls the Select operator.

Returns:: An object of type SelectOutput containing the mapping from nodes to supernodes \(\mathbf{S} \in \mathbb{R}^{N \times K}\).

reduce(**kwargs) → Tensor[source]¶

Calls the Reduce operator.

Returns:: The pooled supernode features \(\mathbf{X}_{\text{pool}}\).

lift(**kwargs) → Tensor | SparseTensor[source]¶

Calls the Lift operator.

Returns:: The un-pooled node features \(\mathbf{X}_{\text{lift}} \approx \mathbf{X}\).

connect(**kwargs) → Tuple[Tensor | SparseTensor, Tensor | None][source]¶

Calls the Connect operator.

Returns:: The adjacency matrix of the coarse graph \(\mathbf{A}_{\text{pool}}\).

preprocessing(x: Tensor, edge_index: Tensor | SparseTensor, **kwargs) → Tuple[Tensor, Tensor | SparseTensor, Tensor | None][source]¶: Preprocess inputs, if needed.

property is_dense: bool¶: Returns True if the pooler uses dense assignments.

property is_sparse: bool¶: Returns True if the pooler uses sparse assignments.

property has_loss: bool¶: Returns True if the pooler overrides compute_loss.

property is_trainable: bool¶: Returns True if any parameter belonging to the pooler or any of its registered sub-modules is trainable.

clear_cache()[source]¶: Clear the caching done by \(\texttt{select}\) and \(\texttt{connect}\).

property is_precoarsenable: bool¶: Returns True if the pooler is precoarsenable.

classmethod get_signature() → Signature[source]¶: Get signature of the pooler’s __init__ function.

classmethod get_forward_signature() → Signature[source]¶: Get signature of the pooler’s forward function.

static data_transforms()[source]¶: Transforms to apply to the dataset before passing it to the model.

class DenseSRCPooling(selector: Select | None = None, reducer: Reduce | None = None, lifter: Lift | None = None, connector: Connect | None = None, cached: bool = False, adj_transpose: bool = False, batched: bool = True, sparse_output: bool = False, cache_preprocessing: bool = False)[source]¶

A base class for dense pooling layers that extends SRCPooling.

It provides a preprocessing function that transform a batch of graphs in sparse representation into a batch of dense graphs. When batched=True, dense poolers accept either raw sparse inputs (which are converted internally) or already-dense padded tensors. In the latter case, an external input-node validity mask can be provided to mark real (non-padded) nodes; otherwise a full-ones mask is assumed.

Parameters:

selector (Select) – The dense \(\texttt{select}\) operator.
reducer (Reduce) – The dense \(\texttt{reduce}\) operator.
lifter (Lift) – The dense \(\texttt{lift}\) operator.
connector (Connect) – The dense \(\texttt{connect}\) operator.
cached (bool, optional) – If set to True, will cache the Select output and the Connect output. (default: False)
cache_preprocessing (bool, optional) – If set to True, will cache the dense adjacency produced by preprocessing(). This should only be enabled when the same graph is reused across iterations (e.g., transductive single-graph tasks). (default: False)
adj_transpose (bool, optional) – If True, the preprocessing step and the DenseConnect operation returns transposed adjacency matrices, so that they could be passed “as is” to the dense message-passing layers. (default: True)
batched (bool, optional) – If True, sparse inputs are converted to batched dense tensors internally. If False, the pooler expects unbatched dense assignments without padding. (default: True)
sparse_output (bool, optional) – If True, the pooled outputs are returned as block-diagonal sparse representations. If False, outputs are returned in batched dense form. (default: False)

preprocessing(x: Tensor, edge_index: Tensor | SparseTensor, edge_weight: Tensor | None = None, batch: Tensor | None = None, max_num_nodes: int | None = None, batch_size: int | None = None, use_cache: bool = False, **kwargs) → Tuple[Tensor, Tensor, Tensor][source]¶

Preprocess inputs for dense pooling methods.

Transform a batch of graphs in sparse representation into a batch of graphs with dense representation.

Parameters:

x (Tensor) – The node features. A tensor of shape \([N, F]\), where \(N\) is the total number of nodes in the batch and \(F\) is the number of node features.
edge_index (Tensor) – The edge indices. A tensor of 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)
max_num_nodes (int, optional) – The maximum number of nodes of a graph in the batch. (default: None)
batch_size (int, optional) – The number of graphs in the batch. (default: None)
use_cache (bool, optional) – If True, it stores the preprocessed adjacency matrix in preprocessing_cache. This is intended for static, single-graph inputs. (default: False)

Returns:

A tuple (x, adj, mask) where x contains batched node features of shape \([B, N_\text{max}, F]\), adj contains batched dense adjacencies of shape \([B, N_\text{max}, N_\text{max}]\), and mask marks real versus padded nodes with shape \([B, N_\text{max}]\).

Return type:

tuple[Tensor, Tensor, Tensor]

clear_cache()[source]¶: Clear the caching done by \(\texttt{select}\) and \(\texttt{connect}\) and the preprocessing cache.

The pooling output of a model of class SRCPooling.

Parameters:

x (Tensor) – The pooled node features.
edge_index (Tensor) – The edge indices of the pooled graph.
edge_weight (Tensor, optional) – The edge features of the coarsened graph. (default: None)
batch (Tensor, optional) – The batch vector of the pooled nodes.
so (SelectOutput) – The selection output. (default: None)
mask – Derived from so.out_mask when so is set. This is a pooled-supernode validity mask of shape \([B, K]\) (or \([1, K]\) for single-graph dense assignments).
loss (Optional[Dict], optional) – The loss dictionary. (default: None)

property mask: Tensor | None¶: Pooled-supernode validity mask, derived from so.out_mask.

property has_loss¶: Returns True if the pooling output has a loss.

get_loss_value(name: str | None = None) → float | List[float][source]¶

Returns the value of the loss with name name or all losses. If the pooling output does not have a loss, it returns 0.

Parameters:: name (str, optional) – The name of the loss to return. If None, returns all losses. (default: None)
Returns:: The value of the loss name or all losses.
Return type:: Union[float, List[float]]

as_data()[source]¶

Converts the pooling output to a Data object.

Returns:: The pooling output as a Data object.
Return type:: Data