Pin Jax to last working release
# What does this PR do?

Jax maintainers has apparently no problem breaking the `optax` library and every library depending on it. So this PR pins Jax and jaxlib for now until this gets sorted out.

cc @patil-suraj @patrickvonplaten for information, will merge as soon as CI is green (or only 500 errors coming from the Hub).
_The documentation is not available anymore as the PR was closed or merged._<|||||>Thank you for fixing this @sgugger !