def from_config(
      cls,
      name: str,
      init_base_rng: jax.Array | None,
      across_axes: dict[str, int],
      epsilon: float | jax.Array = 1e-6,
      dtype: jax.typing.DTypeLike = jnp.float32,
  ) -> LayerNorm:
    """Constructs a layer normalization layer.

    Args:
      name: The name of the layer.
      init_base_rng: The base RNG to use for initializing model 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 ``LayerNorm`` layer.
    """
    return cls([
        Standardize(across=tuple(across_axes.keys()), epsilon=epsilon),
        linear_and_affine.Linear.from_config(
            name=f"{name}/scale",
            init_base_rng=init_base_rng,
            input_axes={},
            output_axes={},
            parallel_axes=across_axes,
            initializer=linear_and_affine.constant_initializer(1.0),
            dtype=dtype,
        ),
        linear_and_affine.AddBias.from_config(
            name=f"{name}/shift",
            init_base_rng=init_base_rng,
            biased_axes=across_axes,
            dtype=dtype,
        ),
    ])
Constructs a layer normalization layer.

    Args:
      name: The name of the layer.
      init_base_rng: The base RNG to use for initializing model 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 ``LayerNorm`` layer.
    
penzai/nn/standardization.py