def __call__(self, value: NamedArray) -> 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 * jnp.reciprocal(jnp.sqrt(var + self.epsilon))

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