3104) because internal tests are breaking.

Commit

2 years ago

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

References

Author

chiamp

Committer

a-googler

Parents