def __call__(self, value: NamedArray) -> NamedArray:
    """Standardizes a named array across the given axes."""
    # Standardize over all of the (local) positional axes in case `across`
    # has more than one axis. All other axes will be vectorized over
    # automatically.
    across = (self.across,) if isinstance(self.across, str) else self.across
    return named_axes.nmap(jax.nn.standardize)(
        value.untag(*across), axis=None, epsilon=self.epsilon
    ).tag(*across)
Standardizes a named array across the given axes.
penzai/deprecated/v1/nn/standardization.py