Fix head pytest breakage

google / flax

Flax is a neural network library for JAX that is designed for flexibility.

https://flax.readthedocs.io

Apache License 2.0

5.78k stars 610 forks source link

Closed IvyZX closed 2 weeks ago

IvyZX commented 2 weeks ago

Removed trivial dependency of tensorflow on jax_utils_test.py
Ignored JAX-side deprecation warning introduced by Tensorflow
No longer skipping jax.tree_.* deprecation warnings because we are free of them!
Pinned orbax-checkpoint to 0.5.15 because 16 and 17 introduces new bugs