Layer Normalization (LayerNorm)#

API#

class warp_nn.modules.layers.LayerNorm(
normalized_shape: int | tuple[int, ...],
*,
eps: float = 1e-05,
elementwise_affine: bool = True,
bias: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply Layer Normalization over the trailing dimensions of the input.

\[\text{LayerNorm}(x) = \frac{x - \text{E}[x]}{\sqrt{\text{Var}[x] + \epsilon}} \, \gamma + \beta\]

The mean and the (biased) variance are computed over the last len(normalized_shape) dimensions.

Learnable parameters:

Name

Shape

Description

\(\gamma\)

weight

normalized_shape

Weights. Only if elementwise_affine is true

\(\beta\)

bias

normalized_shape

Bias. Only if elementwise_affine and bias are true

The weights are initialized to ones and the bias to zeros.

Parameters:

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

  • eps – The value added to the variance for numerical stability.

  • elementwise_affine – Whether to include learnable per-element affine parameters.

  • bias – Whether to include a bias term (only if elementwise_affine is true).

  • 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.