flax

80c91a0b - Unify normalization layers

Unify normalization layers - Refactor normalization helpers to share code and remove subtle differences in implemenation. - Use clipping to avoid negative variances due to fusion optimization.