flax

a7f32a42 - Avoid float32 promotion in BatchNorm when using_running_average=True and force_float32_reductions=False