flax

40f080ed - Avoid using float32 in normalization for mean/var and scale/bias parameters when force_float32_reductions=False