RMSLayerNorm#
- class penzai.deprecated.v1.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 ofSequential
.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.