lgatr.layers.slim_layers.SlimRMSNorm

class lgatr.layers.slim_layers.SlimRMSNorm(v_channels, s_channels, epsilon=0.01, elementwise_affine=True)[source]

Bases: Module

Joint RMS normalization over vector and scalar features.

For vectors the absolute value of the squared norm is used; otherwise the squared norm could be negative under the Lorentz metric.

Parameters:

  • v_channels (int) – Number of vector channels.

  • s_channels (int) – Number of scalar channels.

  • epsilon (float) – Small numerical offset to avoid instabilities.

  • elementwise_affine (bool) – Whether to learn a per-channel gain for the vector and scalar streams.

forward(vectors, scalars)[source]

Normalize jointly.

Parameters:

  • vectors (Tensor) – Lorentz vectors of shape (..., 4, v_channels).

  • scalars (Tensor) – Scalar features of shape (..., s_channels).

Return type:

tuple[Tensor, Tensor]

Returns:

  • outputs_v – Normalized Lorentz vectors, same shape as vectors.

  • outputs_s – Normalized scalar features, same shape as scalars.