Skip to main content

Python module

max.experimental.nn.norm

Normalization layers for MAX neural networks.

Normalization layers

GemmaRMSNormComputes the Root Mean Square normalization on inputs.
GroupNormGroup normalization block.
LayerNormLayer normalization over the last dimension.
RMSNormComputes the Root Mean Square normalization on inputs.

Functions

rms_normApplies Root Mean Square layer normalization to an input tensor.