flax

Rolling back [#3104](https://github.com/google/flax/pull/3104) because internal tests are breaking.

#3126

Merged