Upgrade jax/jaxlib/flax pin versions
# What does this PR do?

So we can have latest TF versions.
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_24791). All of your documentation changes will be reflected on that endpoint.