installing jaxlib through pip

[tbcryoem]

While trying to install jaxlib==0.1.74 using pip, I see the following error. Is there a workaround?

pip install jaxlib==0.1.74
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.74 (from versions: 0.1.60, 0.1.61, 0.1.62, 0.1.63, 0.1.64, 0.1.65, 0.1.66, 0.1.67, 0.1.68, 0.1.69)

Alphafold/Docker latest version requires jaxlib==0.1.74
_check_jaxlib_version
I1210 13:49:26.701840 139826799912768 run_docker.py:222] raise ValueError(msg)
I1210 13:49:26.702075 139826799912768 run_docker.py:222] ValueError: jaxlib is version 0.1.69, but this version of jax requires version 0.1.74.

[jakevdp]

Which Python version are you using? From the "versions available" in your error message, I suspect you're on Python 3.6. Newer versions of JAX & jaxlib require Python 3.7 or later.

[tbcryoem]

# python --version Python 2.7.5 # hostnamectl Operating System: CentOS Linux 7 (Core) CPE OS Name: cpe:/o:centos:centos:7 Kernel: Linux 3.10.0-1160.45.1.el7.x86_64 Architecture: x86-64 # pip install --upgrade pip Requirement already satisfied: pip in /usr/local/lib/python3.6/site-packages (21.3.1) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv # pip install --upgrade jax jaxlib Requirement already satisfied: jax in /usr/local/lib/python3.6/site-packages (0.2.17) Requirement already satisfied: jaxlib in /usr/local/lib/python3.6/site-packages (0.1.69) Requirement already satisfied: numpy>=1.17 in /usr/local/lib64/python3.6/site-packages (from jax) (1.19.5) Requirement already satisfied: absl-py in /usr/local/lib/python3.6/site-packages (from jax) (0.13.0) Requirement already satisfied: opt_einsum in /usr/local/lib/python3.6/site-packages (from jax) (3.3.0) Requirement already satisfied: scipy in /usr/local/lib64/python3.6/site-packages (from jaxlib) (1.5.4) Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.6/site-packages (from jaxlib) (2.0) Requirement already satisfied: six in /usr/local/lib/python3.6/site-packages (from absl-py->jax) (1.16.0) WARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv # pip show jaxlib Name: jaxlib Version: 0.1.69 Summary: XLA library for JAX Home-page: https://github.com/google/jax Author: JAX team Author-email: ***@***.*** License: Apache-2.0 Location: /usr/local/lib/python3.6/site-packages Requires: absl-py, flatbuffers, numpy, scipy Required-by: # pip install jaxlib==0.1.74 ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.74 (from versions: 0.1.60, 0.1.61, 0.1.62, 0.1.63, 0.1.64, 0.1.65, 0.1.66, 0.1.67, 0.1.68, 0.1.69) ERROR: No matching distribution found for jaxlib==0.1.74 # 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 Requirement already satisfied: jax[cuda] in /usr/local/lib/python3.6/site-packages (0.2.17) WARNING: jax 0.2.17 does not provide the extra 'cuda' Requirement already satisfied: numpy>=1.17 in /usr/local/lib64/python3.6/site-packages (from jax[cuda]) (1.19.5) Requirement already satisfied: absl-py in /usr/local/lib/python3.6/site-packages (from jax[cuda]) (0.13.0) Requirement already satisfied: opt_einsum in /usr/local/lib/python3.6/site-packages (from jax[cuda]) (3.3.0) Requirement already satisfied: six in /usr/local/lib/python3.6/site-packages (from absl-py->jax[cuda]) (1.16.0) # nvcc -V nvcc: NVIDIA (R) Cuda compiler driver Copyright (c) 2005-2021 NVIDIA Corporation Built on Sun_Aug_15_21:14:11_PDT_2021 Cuda compilation tools, release 11.4, V11.4.120 Build cuda_11.4.r11.4/compiler.30300941_0 # nvidia-smi Fri Dec 10 18:26:54 2021 +-----------------------------------------------------------------------------+ | NVIDIA-SMI 495.29.05 Driver Version: 495.29.05 CUDA Version: 11.5 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | | | | MIG M. | |===============================+======================+======================| | 0 NVIDIA TITAN X ... Off | 00000000:04:00.0 On | N/A | | 32% 51C P8 14W / 250W | 582MiB / 12192MiB | 0% Default | | | | N/A | +-------------------------------+----------------------+----------------------+ | 1 NVIDIA TITAN X ... Off | 00000000:05:00.0 Off | N/A | | 23% 39C P8 10W / 250W | 2MiB / 12196MiB | 0% Default | | | | N/A | +-------------------------------+----------------------+----------------------+ | 2 NVIDIA TITAN X ... Off | 00000000:08:00.0 Off | N/A | | 23% 35C P8 9W / 250W | 2MiB / 12196MiB | 0% Default | | | | N/A | +-------------------------------+----------------------+----------------------+ | 3 NVIDIA TITAN X ... Off | 00000000:09:00.0 Off | N/A | | 23% 39C P8 11W / 250W | 2MiB / 12196MiB | 0% Default | | | | N/A | +-------------------------------+----------------------+----------------------+ +-----------------------------------------------------------------------------+ | Processes: | | GPU GI CI PID Type Process name GPU Memory | | ID ID Usage | |=============================================================================| | 0 N/A N/A 2749 G /usr/bin/X 346MiB | | 0 N/A N/A 3315 G /usr/bin/gnome-shell 27MiB | | 0 N/A N/A 4544 G /usr/lib64/firefox/firefox 205MiB |

…

________________________________ From: Jake Vanderplas ***@***.***> Sent: Friday, December 10, 2021 4:53 PM To: google/jax ***@***.***> Cc: Biswas, Tapan ***@***.***>; Author ***@***.***> Subject: Re: [google/jax] installing jaxlib through pip (Discussion #8906) Which Python version and operating system are you using? — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub<https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_google_jax_discussions_8906-23discussioncomment-2D1789318&d=DwMCaQ&c=-35OiAkTchMrZOngvJPOeA&r=78b7ZhQK3p1kwpN6M-3oKW1kfRt3GpEihReCBTcSLXw&m=hhMoFOqzmjqdtkvHWATdktcbU-qBZB_16DTxGIkoIXftQ_qRxkPcfQiUk_4R_zZp&s=n5RAnCOtLXg1ZwJds-Nyejy1NOmcLTY7rbEACSGxD7k&e=>, or unsubscribe<https://urldefense.proofpoint.com/v2/url?u=https-3A__github.com_notifications_unsubscribe-2Dauth_APCH65WGXLSAUVPHIZVYWMLUQKOI7ANCNFSM5J2KRQPQ&d=DwMCaQ&c=-35OiAkTchMrZOngvJPOeA&r=78b7ZhQK3p1kwpN6M-3oKW1kfRt3GpEihReCBTcSLXw&m=hhMoFOqzmjqdtkvHWATdktcbU-qBZB_16DTxGIkoIXftQ_qRxkPcfQiUk_4R_zZp&s=Ggk6wY2Z79jI4BoIOr_8HQCPRDLJf80DH_ukFqC2Kuo&e=>. Triage notifications on the go with GitHub Mobile for iOS<https://urldefense.proofpoint.com/v2/url?u=https-3A__apps.apple.com_app_apple-2Dstore_id1477376905-3Fct-3Dnotification-2Demail-26mt-3D8-26pt-3D524675&d=DwMCaQ&c=-35OiAkTchMrZOngvJPOeA&r=78b7ZhQK3p1kwpN6M-3oKW1kfRt3GpEihReCBTcSLXw&m=hhMoFOqzmjqdtkvHWATdktcbU-qBZB_16DTxGIkoIXftQ_qRxkPcfQiUk_4R_zZp&s=iMdR_5edLIkF6SEdUeWfnp27dvHG1_0_2RQ9NAcrxFk&e=> or Android<https://urldefense.proofpoint.com/v2/url?u=https-3A__play.google.com_store_apps_details-3Fid-3Dcom.github.android-26referrer-3Dutm-5Fcampaign-253Dnotification-2Demail-2526utm-5Fmedium-253Demail-2526utm-5Fsource-253Dgithub&d=DwMCaQ&c=-35OiAkTchMrZOngvJPOeA&r=78b7ZhQK3p1kwpN6M-3oKW1kfRt3GpEihReCBTcSLXw&m=hhMoFOqzmjqdtkvHWATdktcbU-qBZB_16DTxGIkoIXftQ_qRxkPcfQiUk_4R_zZp&s=03XSkQASKdyG7AQKdEwf9KfH7A0MzDjuOFTJ_lwHmno&e=>.

[jakevdp]

If you’re on Python 2.7, I wouldn’t expect JAX to work at all. JAX requires Python 3.7 or newer. I’d start by ensuring you’re using a newer Python, and then try again.

[tbcryoem]

Thanks, Upon updating to Python 3.7.9, installation of a specific jaxlib version works; # python3 --version Python 3.7.9 # pip3 --version pip 21.3.1 from /usr/local/lib/python3.7/site-packages/pip (python 3.7) # pip show jaxlib Name: jaxlib Version: 0.1.75 Summary: XLA library for JAX Home-page: https://github.com/google/jax Author: JAX team Author-email: jax-dev@google.com License: Apache-2.0 Location: /usr/local/lib/python3.7/site-packages Requires: absl-py, flatbuffers, numpy, scipy Required-by: # pip uninstall jaxlib Found existing installation: jaxlib 0.1.75 Uninstalling jaxlib-0.1.75: Would remove: /usr/local/lib/python3.7/site-packages/jaxlib-0.1.75.dist-info/* /usr/local/lib/python3.7/site-packages/jaxlib/* Proceed (Y/n)? Y Successfully uninstalled jaxlib-0.1.75 # pip install jaxlib==0.1.74 Collecting jaxlib==0.1.74 Downloading jaxlib-0.1.74-cp37-none-manylinux2010_x86_64.whl (61.5 MB) |████████████████████████████████| 61.5 MB 15.5 MB/s Requirement already satisfied: numpy>=1.18 in /usr/local/lib/python3.7/site-packages (from jaxlib==0.1.74) (1.21.4) Requirement already satisfied: absl-py in /usr/local/lib/python3.7/site-packages (from jaxlib==0.1.74) (0.13.0) Requirement already satisfied: flatbuffers<3.0,>=1.12 in /usr/local/lib/python3.7/site-packages (from jaxlib==0.1.74) (2.0) Requirement already satisfied: scipy in /usr/local/lib/python3.7/site-packages (from jaxlib==0.1.74) (1.7.3) Requirement already satisfied: six in /usr/local/lib/python3.7/site-packages (from absl-py->jaxlib==0.1.74) (1.16.0) Installing collected packages: jaxlib Successfully installed jaxlib-0.1.74 XXXXXXX Unrelated problem: (although not sure yet why docker does not sense the new version for running alphafold) run_docker.py:222] ValueError: jaxlib is version 0.1.69, but this version of jax requires version 0.1.74

[hawkinsp]

I think the Alphafold docker image pins a particular known working-for-Alphafold version of Jax. So it might be easiest just to use that version! That's the point of the docker image after all.