Ctrl+K

RMSLayerNorm

RMSLayerNorm#

class penzai.nn.standardization.RMSLayerNorm[source]#

Bases: Sequential

Root-mean-squared layer normalization layer.

RMS layer normalization layers consist of:

  • root-mean-squared standardization over a feature axis or axes,

  • with a learned parallel rescaling of each feature along those axes.

As proposed by Zhang & Sennrich (2019): https://arxiv.org/abs/1910.07467.

For flexibility, RMSLayerNorm is a subclass of Sequential.

Methods

__init__(sublayers)

from_config(across_axes[, epsilon, dtype])

Constructs a RMS layer normalization layer.

Attributes

sublayers

Inherited Methods

(expand to view inherited methods)

attributes_dict()

Constructs a dictionary with all of the fields in the class.

from_attributes(**field_values)

Directly instantiates a struct given all of its fields.

input_structure()

Returns the input structure of this layer.

key_for_field(field_name)

Generates a JAX PyTree key for a given field name.

output_structure()

Returns the output structure of this layer.

select()

Wraps this struct in a selection, enabling functional-style mutations.

tree_flatten()

Flattens this tree node.

tree_flatten_with_keys()

Flattens this tree node with keys.

tree_unflatten(aux_data, children)

Unflattens this tree node.

treescope_color()

__call__(value)

Runs each of the sublayers in sequence.
classmethod from_config(across_axes: dict[str, int], epsilon: float | jax.Array = 1e-06, dtype: jax.typing.DTypeLike = <class 'jax.numpy.float32'>) → RMSLayerNorm[source]#

Constructs a RMS layer normalization layer.

Parameters:

  • across_axes – Names and lengths of the axes to normalize over.

  • epsilon – Epsilon parameter for the standardization step.

  • dtype – Dtype of the scale and shift parameters.

Returns:

A newly-constructed RMSLayerNorm layer.

Contents