lgatr.layers.layer_norm.EquiLayerNorm

class lgatr.layers.layer_norm.EquiLayerNorm(mv_channels=0, s_channels=0, mv_channel_dim=-2, epsilon=0.01, gain=1.0, elementwise_affine=False)[source]

Bases: Module

Layer normalization for multivectors (and an optional scalar stream).

Rescales the multivector input such that mean_channels |x|^2 = 1, where the norm is the GA norm and the mean is taken over the channel dimension. The scalar stream, if present, is passed through a regular torch.nn.functional.layer_norm().

With elementwise_affine=True a learnable gain is applied after normalization: a per-grade scalar per multivector channel (shape (mv_channels, 5), broadcast over the components of each grade) and a per-channel scalar for the scalar stream (shape (s_channels,)). Scaling each grade independently preserves Pin-equivariance.

Parameters:

  • mv_channels (int) – Number of multivector channels. Only used to size the affine gain when elementwise_affine=True.

  • s_channels (int) – Number of scalar channels. Only used to size the affine gain when elementwise_affine=True.

  • mv_channel_dim (int) – Channel-dimension index for multivector inputs. Defaults to the second-to-last entry (the last is the multivector component dimension).

  • epsilon (float) – Small numerical offset to avoid instabilities. The default is intentionally larger than usual to balance the fact that some multivector components do not contribute to the norm.

  • elementwise_affine (bool) – Whether to learn a per-channel-per-grade multivector gain and a per-channel scalar gain.

forward(multivectors, scalars=None)[source]

Apply equivariant LayerNorm.

Parameters:

  • multivectors (Tensor) – Multivector inputs of shape (..., channels, 16).

  • scalars (Tensor | None) – Optional scalar inputs of shape (..., s_channels). If None, no scalar normalization is performed and outputs_s is None.

Return type:

tuple[Tensor, Tensor | None]

Returns:

  • outputs_mv – Normalized multivectors of shape (..., channels, 16).

  • outputs_s – Normalized scalars of shape (..., s_channels), or None if scalars is None.