IrrepSubspaceNormPooling

class IrrepSubspaceNormPooling(in_rep)

Bases: Module

Pool irrep features into \(\mathbb{G}\)-invariant radii.

Given \(\mathbf{x}\in\mathcal{X}\) with representation \(\rho_{\mathcal{X}}\), the module computes one scalar per irreducible copy in the isotypic/irrep-spectral basis:

\[r_{k,i} = \lVert \hat{\mathbf{x}}_{k,i} \rVert_2, \qquad \hat{\mathbf{x}}=\mathbf{Q}^T\mathbf{x}.\]

This is exactly irrep_radii(), exposed as a module. The output transforms under a direct sum of trivial representations, hence is invariant.

Parameters:

in_rep (Representation) – Representation \(\rho_{\text{in}}\) describing how the input last dimension transforms.

forward(x)

Compute one invariant radius per irreducible copy.

Parameters:

x (Tensor) – Input with trailing dimension in_rep.size; any leading batch/time dims are accepted.

Returns:

Tensor with same leading shape as x and last dim out_rep.size containing one

Euclidean norm per irrep block (trivial features).

Return type:

Tensor