eBatchNorm1d#

class eBatchNorm1d(in_rep, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)[source]#

Bases: Module

Symmetry-aware Batch Normalization over the representation dimension.

The mean and variance are computed with var_mean(), enforcing that each irreducible subspace shares a single variance scalar. The optional affine parameters are implemented via eAffine to preserve equivariance.

The layer satisfies:

\[\rho_{\text{in}}(g) \mathbf{y} = \text{BatchNorm}(\rho_{\text{in}}(g) \mathbf{x})\]
Parameters:

  • in_rep (Representation) – Representation \(\rho_{\text{in}}\) describing the feature space.

  • eps (float) – Numerical stabilizer added to the variance.

  • momentum (float) – Momentum for exponential moving averages.

  • affine (bool) – If True, apply a symmetry-preserving affine transform.

  • track_running_stats (bool) – If True, keep running mean/variance buffers.

Shape:

  • Input: (..., in_rep.size)

  • Output: same shape

check_equivariance(atol=1e-05, rtol=1e-05)[source]#

Check equivariance using random group elements.

forward(x)[source]#

Normalize using symmetry-constrained batch statistics and optional equivariant affine map.

Return type:

Tensor

Parameters:

x (Tensor)