Lift

BaseLift

A base class to lift the features of the supernodes back into the original node space.

EigenPoolLift

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

Lift

A template class for the lift operator.

class BaseLift(matrix_op: Literal['transpose', 'inverse', 'precomputed'] = 'precomputed', reduce_op: str = 'sum')[source]

A base class to lift the features of the supernodes back into the original node space.

The lift operation is implemented as

\[\mathbf{X}_{\text{lift}} = f(\mathbf{S}_{\text{inv}}, \mathbf{X}_{\text{pool}}) \approx \mathbf{X}.\]

where

  • \(\mathbf{X}_{\text{lift}} \in \mathbb{R}^{N \times F}\) are the lifted node features,

  • \(\mathbf{S}_{\text{inv}} \in \mathbb{R}^{K \times N}\) is the inverse node assignment operator,

  • \(\mathbf{X}_{\text{pool}} \in \mathbb{R}^{K \times F}\) are the pooled features of the supernodes,

  • \(f(\cdot)\) is the lifting operation that specifies how \(\mathbf{S}_{\text{inv}}\) is used to compute the lifted features \(\mathbf{X}_{\text{lift}}\). In most cases, \(f(\mathbf{S}_{\text{inv}}, \mathbf{X}_{\text{pool}}) = \mathbf{S}_{\text{inv}}^{\top} \mathbf{X}_{\text{pool}}\).

It also works for dense pooling operators. In that case, \(\mathbf{X}_{\text{lift}} \in \mathbb{R}^{B \times N \times F}\), \(\mathbf{S}_{\text{inv}} \in \mathbb{R}^{B \times K \times N}\), \(\mathbf{X}_{\text{pool}} \in \mathbb{R}^{B \times K \times F}\).

Parameters:
  • matrix_op (LiftType) –

    Defines how to compute the matrix \(\mathbf{S}_\text{inv}\) to lift the pooled node features.

    • "precomputed" (default): Use as \(\mathbf{S}_\text{inv}\) what is already stored in the "s_inv" attribute of the SelectOutput.

    • "transpose": Recomputes \(\mathbf{S}_\text{inv}\) as \(\mathbf{S}^\top\), the transpose of \(\mathbf{S}\).

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

  • reduce_op (ReduceType) – The aggregation function to be applied to the lifted node features. Can be any string of class ReduceType admitted by scatter, e.g., 'sum', 'mean', 'max') (default: "sum")

forward(x_pool: Tensor, so: SelectOutput | None = None, batch: Tensor | None = None, batch_pooled: Tensor | None = None, **kwargs) Tensor[source]

Forward pass of the Lift operation.

Parameters:
  • x_pool (Tensor) – The pooled node features.

  • so (SelectOutput) – The output of the \(\texttt{select}\) operator.

  • batch (Tensor, optional) – The batch vector for the original nodes. If not provided, so.batch is used when available. (default: None)

  • batch_pooled (Tensor, optional) – The batch vector for the pooled nodes. For dense multi-graph lifting with flattened pooled features \([B \cdot K, F]\), if not provided it is inferred as contiguous graph blocks of size \(K\). (default: None)

Returns:

The lifted node features.

Return type:

Tensor

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

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

It uses the pooling matrix \(\boldsymbol{\Theta}\) stored in so.theta and lifts pooled features back to node space as:

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

where \(\mathbf{X}_{\text{pool,raw}}\) is the mode-major version of the pooled features.

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 Lift. (default: "sum")

forward(x_pool: Tensor, so: SelectOutput | None = None, batch: Tensor | None = None, batch_pooled: Tensor | None = None, edge_index: Tensor | None = None, edge_weight: Tensor | None = None, **kwargs) Tensor[source]

Forward pass.

Parameters:
  • x_pool (Tensor) – Pooled feature matrix of shape \([K, H\cdot F]\), \([B\cdot K, H\cdot F]\), or \([B, K, H\cdot F]\).

  • so (SelectOutput, optional) – Output of the \(\texttt{select}\) operator with dense assignment matrix so.s and pooling matrix so.theta.

  • batch (Tensor, optional) – Batch vector for original nodes. If None, this method uses so.batch when available. (default: None)

  • batch_pooled (Tensor, optional) – Batch vector for pooled nodes in multi-graph lifting. (default: None)

  • edge_index (Tensor, optional) – Unused by EigenPooling. (default: None)

  • edge_weight (Tensor, optional) – Unused by EigenPooling. (default: None)

Returns:

Lifted node features of shape \([N, F]\).

Return type:

Tensor

class Lift(*args, **kwargs)[source]

A template class for the lift operator.

forward(x_pool: Tensor, so: SelectOutput, **kwargs) Tensor[source]

Forward pass.

Parameters:
  • x_pool (Tensor) – the pooled node features \(\mathbf{X}_{\text{pool}} \in \mathbb{R}^{K \times F}\)

  • so (SelectOutput) – The output of the \(\texttt{select}\) operator.