Root Mean Square Normalization (RMSNorm)#

API#

class warp_nn.modules.layers.RMSNorm(
normalized_shape: int | tuple[int, ...],
*,
eps: float | None = None,
elementwise_affine: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply Root Mean Square (RMS) Normalization over the trailing dimensions of the input.

\[\text{RMSNorm}(x) = \frac{x}{\sqrt{\text{E}[x^2] + \epsilon}} \, \gamma\]

The mean square is computed over the last len(normalized_shape) dimensions.

Learnable parameters:

Name

Shape

Description

\(\gamma\)

weight

normalized_shape

Weights. Only if elementwise_affine is true

The weights are initialized to ones.

Parameters:

  • normalized_shape – The shape of the trailing dimensions to normalize over.

  • eps – The value added to the mean square for numerical stability. If None, the machine epsilon of the 32-bit floating-point type is used.

  • elementwise_affine – Whether to include learnable per-element weights.

  • initialize_parameters – Whether to initialize the parameters with their default/initial values. If false, the parameters are left as uninitialized memory, and reading them (e.g. during a forward pass) before loading their values (e.g. from a state dictionary) is undefined behavior.

  • requires_grad – Whether the parameters and the cached output arrays of the module require gradients.

__call__(
input: array,
) → array[source]#

Forward pass of the module.

Parameters:

input – The input array, with shape (*, *normalized_shape).

Returns:

The output array, with same shape as the input array.

Raises:

ValueError – If the trailing dimensions of the input array do not match the normalized shape.