flax

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

#4314

Merged