Flax requires jax>=0.2.21 ?

[jrabary]

Dear all,

I'm trying to install Flax but it requires a version of jax>=0.2.21. The problem is that

pip install jax[cuda] only install jax==0.2.17

Is there a quick workaround so we can install flax without waiting for the update of the jax pip package ?

Jao

[jheek]

As per the JAX changelog: could you try the updated install command:

pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

[jrabary]

Here is the output

pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
Looking in links: https://storage.googleapis.com/jax-releases/jax_releases.html
Processing /root/.cache/pip/wheels/c0/04/e9/96543974a96ee70858fe2d8e36dbe24dcf77b32f206ec01112/jax-0.2.17-py3-none-any.whl
  WARNING: jax 0.2.17 does not provide the extra 'cuda'
Requirement already satisfied, skipping upgrade: opt-einsum in /usr/local/lib/python3.6/dist-packages (from jax[cuda]) (3.3.0)
Requirement already satisfied, skipping upgrade: absl-py in /usr/local/lib/python3.6/dist-packages (from jax[cuda]) (0.13.0)
Requirement already satisfied, skipping upgrade: numpy>=1.17 in /usr/local/lib/python3.6/dist-packages (from jax[cuda]) (1.19.5)
Requirement already satisfied, skipping upgrade: six in /usr/local/lib/python3.6/dist-packages (from absl-py->jax[cuda]) (1.15.0)
Installing collected packages: jax
Successfully installed jax-0.2.17

It installs jax==0.2.17 and there's a warning about cuda.

[jheek]

Strange, I can't reproduce this. I think this is a question best asked in the JAX repository

[jrabary]

The problem was the python version jax-ml/jax#8399.