Normalization

The normalization module provides custom normalization layers used throughout the network. These layers are designed to be flexible and efficient, supporting various normalization strategies.

Usage

These normalization layers can be used in two ways:

1. Directly in model implementations

2. Through the layer kernels configuration system

Example configuration using layer kernels:

layer_kernels:
  processor:
    LayerNorm:
      _target_: anemoi.models.layers.normalization.AutocastLayerNorm
      bias: False

The normalization layers are particularly useful when:

1. Working with mixed precision training

2. Implementing ensemble models with noise injection

3. Requiring specialized normalization behavior in specific parts of the model

Available Layers

class anemoi.models.layers.normalization.AutocastLayerNorm(*args, **kwargs)

Bases: LayerNorm

LayerNorm that casts the output back to the input type.

forward(x: Tensor) → Tensor

Forward with explicit autocast back to the input type.

This casts the output to (b)float16 (instead of float32) when we run in mixed precision.

class anemoi.models.layers.normalization.ConditionalLayerNorm(normalized_shape: int | list | Size, condition_shape: int = 16, zero_init: bool = True, autocast: bool = True)

Bases: Module

Conditional Layer Normalization.

x_norm = a(u) * (x - mean) / sqrt(var + eps) + b(u)

forward(x: Tensor, cond: Tensor) → Tensor

Conditional Layer Normalization.

Parameters:

• x (Tensor) – Input tensor to be normalized.

• cond (Tensor) – Conditioning tensor used to modulate the normalization.

Returns:

Output tensor.

Return type:

Tensor