Lift¶
A base class to lift the features of the supernodes back into the original node space. |
|
The \(\texttt{lift}\) operator for EigenPooling. |
|
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 theSelectOutput."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
ReduceTypeadmitted byscatter, 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.batchis 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:
- 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.thetaand 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.sand pooling matrixso.theta.batch (Tensor, optional) – Batch vector for original nodes. If
None, this method usesso.batchwhen 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:
- 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.