Flax requires jax>=0.2.21 but unable to install it

[jrabary]

Dear all,

I've already asked this question on the Flax repository : google/flax#1651.

When installing jax with pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html I always get the version 0.2.17 of jax but flax need the version 0.2.21.

Is there a quick workaround for this problem ?

Best regards,

Jao

[hawkinsp]

I suspect the reason is that you are using Python 3.6. We dropped support for Python 3.6 in jax v0.2.18, per the deprecation policy.

Please update to a supported version of Python.

[jrabary]

Indeed my python version is 3.6.9. Thanks!