flax

3e789b71 - Update references to JAX's GitHub repo

Update references to JAX's GitHub repo JAX has moved from https://github.com/google/jax to https://github.com/jax-ml/jax PiperOrigin-RevId: 703061304