Reduce¶

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

For graph-level aggregation, AggrReduce and GlobalReduce rely on PyG’s torch_geometric.nn.aggr.Aggregation modules. The helper function get_aggr() lets you obtain these aggregators by string alias, mirroring how tgp.poolers.get_pooler() instantiates poolers by alias.

from tgp.reduce import get_aggr, AggrReduce, GlobalReduce

# Aggregator instance from alias
aggr = get_aggr("mean")
reducer = AggrReduce(aggr)

# Graph-level readout with GlobalReduce
readout_layer = GlobalReduce(
    reduce_op="set2set", in_channels=64, processing_steps=3
)
x_graph = readout_layer(x, batch=batch)

Aggregator aliases¶

The table below lists the supported aliases for get_aggr() and the corresponding PyG aggregation classes:

Alias	PyG aggregation class
`sum`	`torch_geometric.nn.aggr.SumAggregation`
`mean`	`torch_geometric.nn.aggr.MeanAggregation`
`max`	`torch_geometric.nn.aggr.MaxAggregation`
`min`	`torch_geometric.nn.aggr.MinAggregation`
`mul`	`torch_geometric.nn.aggr.MulAggregation`
`var`	`torch_geometric.nn.aggr.VarAggregation`
`std`	`torch_geometric.nn.aggr.StdAggregation`
`softmax`	`torch_geometric.nn.aggr.SoftmaxAggregation`
`power_mean`	`torch_geometric.nn.aggr.PowerMeanAggregation`
`median`	`torch_geometric.nn.aggr.MedianAggregation`
`quantile`	`torch_geometric.nn.aggr.QuantileAggregation`
`lstm`	`torch_geometric.nn.aggr.LSTMAggregation`
`gru`	`torch_geometric.nn.aggr.GRUAggregation`
`set2set`	`torch_geometric.nn.aggr.Set2Set`
`degree_scaler`	`torch_geometric.nn.aggr.DegreeScalerAggregation`
`sort`	`torch_geometric.nn.aggr.SortAggregation`
`multi`	`torch_geometric.nn.aggr.MultiAggregation`
`attentional`	`torch_geometric.nn.aggr.AttentionalAggregation`
`equilibrium`	`torch_geometric.nn.aggr.EquilibriumAggregation`
`mlp`	`torch_geometric.nn.aggr.MLPAggregation`
`deep_sets`	`torch_geometric.nn.aggr.DeepSetsAggregation`
`set_transformer`	`torch_geometric.nn.aggr.SetTransformerAggregation`
`lcm`	`torch_geometric.nn.aggr.LCMAggregation`
`variance_preserving`	`torch_geometric.nn.aggr.VariancePreservingAggregation`
`patch_transformer`	`torch_geometric.nn.aggr.PatchTransformerAggregation`
`graph_multiset_transformer`	`torch_geometric.nn.aggr.GraphMultisetTransformer`

`AggrReduce`	Reduce operator that wraps a PyG `torch_geometric.nn.aggr.Aggregation`.
`Reduce`	A template class for implementing the \(\texttt{reduce}\) operator.
`BaseReduce`	The basic \(\texttt{reduce}\) operator that computes \(\mathbf{S}^\top \mathbf{X}\).
`EigenPoolReduce`	The \(\texttt{reduce}\) operator for EigenPooling.
`GlobalReduce`	Graph-level readout as a module wrapping `AggrReduce`.

class AggrReduce(aggr)[source]¶

Reduce operator that wraps a PyG torch_geometric.nn.aggr.Aggregation.

Aggregates node features within each supernode using the given aggregation module. Supports sparse assignment matrices and graph-level readout mode (so=None). Dense SelectOutput assignments are not supported: use BaseReduce for dense/soft reductions.

In SelectOutput, pooled-node indices are named cluster_index for historical reasons; here this is equivalent to “supernode index”.

so=None is supported for graph-level readout: all nodes are assigned to one supernode per graph (using batch as supernode index), or to a single supernode when batch is None.

Parameters:: aggr – A PyG Aggregation instance (e.g. torch_geometric.nn.aggr.SumAggregation, torch_geometric.nn.aggr.MeanAggregation).

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

Aggregate node features according to the supernode assignments.

Parameters:

x (Tensor) – Node features of shape \([N, F]\) or \([B, N, F]\).
so (SelectOutput, optional) – Select output containing assignment information. If None, performs graph-level readout using batch.
batch (Tensor, optional) – Batch vector assigning each node to a graph.
size (int, optional) – Expected number of pooled nodes (readout groups). If None, inferred from so or batch.

Returns:

A pair (x_pool, batch_pool) with pooled features and pooled batch indices.

Return type:

tuple

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

A template class for implementing the \(\texttt{reduce}\) operator.

static reduce_batch(select_output: SelectOutput, batch: Tensor | None) → Tensor | None[source]¶

Computes the batch vector of the coarsened graph.

Parameters:

select_output (SelectOutput) – The output of Select.
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 pooled batch.

Return type:

Tensor or None

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

Forward pass.

Parameters:

x (Tensor) – The node feature matrix. For a sparse pooler, x has shape \([N, F]\), where \(N\) is the number of nodes in the batch and \(F\) is the number of node features. For a dense pooler, x has shape \([B, N, F]\), where \(B\) is the batch size.
so (SelectOutput) – The output of the \(\texttt{select}\) operator.
batch (torch.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 pooled features \(\mathbf{X}_{pool}\) of the supernodes.

Return type:

Tensor

class BaseReduce[source]¶

The basic \(\texttt{reduce}\) operator that computes \(\mathbf{S}^\top \mathbf{X}\).

For sparse assignment \(\mathbf{S}\), this is implemented as a sum over nodes in each cluster (with optional weighting by so.s.values()). For dense assignment, it is a matrix multiply \(\mathbf{S}^\top \mathbf{X}\).

For dense multi-graph batches (dense \([N, K]\) with a batch vector), each graph is processed separately (unbatch then per-graph matmul) for memory efficiency when using unbatched dense poolers (batched=False).

For dense unbatched assignments \([N, K]\) with multi-graph batches, return_batched=True returns \([B, K, F]\); otherwise \([B \cdot K, F]\). For dense batched assignments \([B, N, K]\), output is always \([B, K, F]\).

forward(x: Tensor, so: SelectOutput, *, batch: Tensor | None = None, return_batched: bool = False, **kwargs) → Tuple[Tensor, Tensor | None][source]¶

Forward pass computing \(\mathbf{S}^\top \mathbf{X}\).

Parameters:

x (Tensor) – The node feature matrix. For a sparse pooler, x has shape \([N, F]\), where \(N\) is the number of nodes in the batch and \(F\) is the number of node features. For a dense pooler, x has shape \([B, N, F]\), where \(B\) is the batch size.
so (SelectOutput) – The output of the \(\texttt{select}\) operator.
batch (torch.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)
return_batched (bool, optional) – For dense unbatched \([N, K]\) assignments with multi-graph batches, controls output shape: True gives \([B, K, F]\), False gives \([B \cdot K, F]\). For single-graph \([N, K]\), True wraps output as \([1, K, F]\). (default: False)

class EigenPoolReduce(num_modes: int = 5, reduce_op: str = 'sum')[source]¶

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

It uses the pooling matrix \(\boldsymbol{\Theta}\) computed by EigenPoolSelect and stored in so.theta. For each graph:

\[\mathbf{X}_{\text{pool,raw}} = \boldsymbol{\Theta}^{\top}\mathbf{X},\]

then \(\mathbf{X}_{\text{pool,raw}}\) is reshaped from mode-major layout \([H \cdot K, F]\) to \([K, H \cdot F]\).

return_batched is only used when returning multi-graph results (stack vs concatenate). EigenPool uses 2D inputs with a batch vector; there is no separate dense path.

Parameters:

num_modes (int, optional) – Number of eigenvector modes \(H\). Kept for API symmetry with the EigenPool components. (default: 5)
reduce_op (ReduceType, optional) – Kept for API compatibility with Reduce. (default: "sum")

forward(x: Tensor, so: SelectOutput, *, batch: Tensor | None = None, edge_index: Tensor | None = None, edge_weight: Tensor | None = None, return_batched: bool = False, **kwargs) → Tuple[Tensor, Tensor | None][source]¶

Forward pass.

Parameters:

x (Tensor) – Node feature matrix \(\mathbf{X}\) of shape \([N, F]\).
so (SelectOutput) – Output of the \(\texttt{select}\) operator with dense assignment matrix so.s and pooling matrix so.theta.
batch (Tensor, optional) – Batch vector for sparse multi-graph inputs. If None, this method uses so.batch when available. (default: None)
edge_index (Tensor, optional) – Unused by EigenPooling. (default: None)
edge_weight (Tensor, optional) – Unused by EigenPooling. (default: None)
return_batched (bool, optional) – Only used for multi-graph batches. If True, returns \([B, K, H \cdot F]\); otherwise concatenated \([B \cdot K, H \cdot F]\). For single graphs, if True returns \([1, K, H \cdot F]\). (default: False)

Returns:

Pooled features and pooled batch vector.

Return type:

(~torch.Tensor, ~torch.Tensor or None)

class GlobalReduce(reduce_op: str | object = 'sum', **aggr_kwargs)[source]¶

Graph-level readout as a module wrapping AggrReduce.

This module aggregates node features to one vector per graph, supporting both sparse inputs [N, F] with an optional batch vector and dense inputs [B, N, F] with an optional boolean mask of shape [B, N].

Parameters:

reduce_op – Aggregation to use: a string alias (e.g. \"sum\", \"mean\", \"max\", \"min\", \"lstm\", \"set2set\") or a PyG torch_geometric.nn.aggr.Aggregation instance. Strings are resolved via get_aggr().
**aggr_kwargs – Passed to get_aggr() when reduce_op is a string (e.g. in_channels, out_channels, processing_steps).

forward(x: Tensor, batch: Tensor | None = None, size: int | None = None, mask: Tensor | None = None) → Tensor[source]¶

Aggregate node features to one vector per graph.

Infers sparse vs dense from x.ndim: 2D [N, F] is sparse (use batch for grouping); 3D [B, N, F] is dense (reduce over the node dimension). Nodes must be on the second-to-last dimension.

Parameters:

x – Node features. Shape [N, F] (sparse) or [B, N, F] (dense).
batch – Batch vector for sparse x, shape [N]. Ignored for dense.
size – Number of graphs for sparse readout when batch is provided. Passing size with sparse x and batch=None raises ValueError.
mask – Input-node validity mask for batched (dense) x only, shape [B, N]. Passing mask with sparse x or with a mismatched shape raises ValueError.

Returns:

Tensor of shape [B, F] (or [1, F] for single graph sparse).

get_aggr(alias: str, **kwargs: Any) → Any[source]¶

Return a PyG torch_geometric.nn.aggr.Aggregation instance by alias.

This is the aggregation analogue of tgp.poolers.get_pooler(): instead of instantiating poolers by alias, get_aggr() instantiates PyG aggregation modules. Use it with AggrReduce or GlobalReduce when you prefer to configure aggregation via a string name.

Parameters:

alias – Name of the aggregator (e.g. "sum", "mean", "lstm", "set2set"). Case-insensitive; dashes are normalized to underscores. The full list of supported aliases is documented in tgp.reduce (see the Aggregator aliases section).
**kwargs – Passed to the aggregator constructor. Parametrized aggregators typically need in_channels, out_channels, and/or processing_steps (e.g. for Set2Set, LSTM). Any extra keyword arguments that are not accepted by the underlying PyG class are silently dropped.

Returns:

An instance of the requested PyG Aggregation.

Raises:

ImportError – If torch_geometric.nn.aggr is not available.
ValueError – If alias is not recognized.

Example

>>> from tgp.reduce import get_aggr, AggrReduce, GlobalReduce
>>> # Inside a pooling layer
>>> reducer = AggrReduce(get_aggr("mean"))
>>> # For graph-level readout with Set2Set
>>> readout_layer = GlobalReduce(
...     reduce_op="set2set", in_channels=64, processing_steps=3
... )
>>> x_graph = readout_layer(x, batch=batch)