flax

e6379c1b - option of forcing the input of softmax to be fp32 for better numerical stability in mixed-precision training.