flax

f400015e - Copybara import of the project:

Copybara import of the project: -- ff8524813e1bf7aa75dee17495ffb0d411065590 by IvyZX <ivyzhengx@gmail.com>: Correctly expose flax.config.temp_flip_flag PiperOrigin-RevId: 811006892