lgatr.primitives.normalization

Multivector normalization.

Functions

equi_layer_norm(x[, channel_dim, gain, epsilon])

Equivariant LayerNorm for multivectors.
lgatr.primitives.normalization.equi_layer_norm(x, channel_dim=-2, gain=1.0, epsilon=0.01)[source]

Equivariant LayerNorm for multivectors.

Rescales the input such that mean_channels |x|^2 = 1, where the norm is the GA norm and the mean is taken over the channel dimension.

Using a factor gain > 1 makes up for the fact that the GA norm overestimates the actual standard deviation of the input data.

Parameters:

  • x (Tensor) – Input multivectors of shape (..., channels, 16).

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

  • gain (float | Tensor) – Target output scale.

  • epsilon (float | Tensor) – 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.

Returns:

Normalized multivectors of shape (..., channels, 16).

Return type:

outputs