def __call__(self, value: NamedArray, **_unused_side_inputs) -> NamedArray:
    """Root-mean-square standardizes a named array across the given axes."""
    across = (self.across,) if isinstance(self.across, str) else self.across

    @named_axes.nmap
    def _rms_standardize(x):
      var = jnp.mean(jnp.square(x))
      return x * jax.lax.rsqrt(var + self.epsilon)

    return _rms_standardize(value.untag(*across)).tag(*across)
Root-mean-square standardizes a named array across the given axes.
penzai/nn/standardization.py