Un-pin JAX from <= 0.3.6

[sanchit-gandhi]

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

[sgugger]

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.

[sanchit-gandhi]

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.

[sanchit-gandhi]

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[sanchit-gandhi]

Closed by https://github.com/huggingface/transformers/pull/24791