Un-pin JAX from <= 0.3.6
### Feature request

JAX was pinned to <= 0.3.6 in #16808 when a minor release of JAX was published that broke Optax:
https://github.com/huggingface/transformers/blob/8b2501b4b9c5afad0dca2c964bc31b9fca09df4e/setup.py#L123
There have been 18 subsequent minor releases since (_c.f._ https://jax.readthedocs.io/en/latest/changelog.html), with v0.3.24 the latest. We should un-pin this hard requirement to allow users to install the latest version of JAX with Transformers. 

### Motivation

Should a user have JAX == 0.3.24 installed and then tries to install Transformers using pip:
```
pip install transformers[flax]
```
JAX is downgraded to 0.3.6 due to the pinning requirement. Should a user be using the latest JAX features, this then requires JAX to be _re-upgraded_ to 0.3.24:
```
pip install -U jax
```
The same holds true for JAX dependencies (e.g. Flax or Optax).

### Your contribution

Investigate all versions of JAX/Flax/Optax triplets compatible with Transformers. Black-list those that break the CI, otherwise allow users to install that version of JAX in `setup.py`.

cc @cgarciae
Until the Jax/Flax team guarantees they will maintain compatibility in their ecosystem during releases, they should be an upper pin to avoid sudden breaks of the CI.<|||||>Very much agree @sgugger! I'll maintain an upper-bound on the versions of JAX and JAX derived libraries to ensure there aren't minor releases that break our CI 🤗 We can then deal with new versions of JAX on a version-by-version basis.<|||||>[](url)