Not able to install jaxlib

Admin 1
Viewed 20702

I am trying to install jaxlib on my windows 10 by the following command which I found on the documentation..

pip install jaxlib

It shows the following error

Collecting jaxlib
  Could not find a version that satisfies the requirement jaxlib (from versions: None)
No matching distribution found for jaxlib
jax pip installation numpy python
5 Answers
Admin1

I went through the process the last two days myself so here is what i did:

  1. download and install the latest version of microsoft visual studio ( to get a c++ compiler and toolchain)

  2. download and install python

  3. create a virtual python environment with the packages pip, setuptools, six, wheel and numpy (i did it in the GUI of pycharm)

  4. download jax

  1. open up a windows powershell as administrator, change to the jax directory and complete the following steps (commands are in quotes)
  • install chocolatey (package manager for easy bazel installation)

Set-ExecutionPolicy Bypass -Scope Process -Force; [System.Net.ServicePointManager]::SecurityProtocol = [System.Net.ServicePointManager]::SecurityProtocol -bor 3072; iex ((New-Object System.Net.WebClient).DownloadString('https://community.chocolatey.org/install.ps1'))

  • install bazel (for building jaxlib)

choco install bazel

  • install msys2 (linux utilities for bazel)

choco install msys2

  • permamently link the python environment in your powershell

[System.Environment]::SetEnvironmentVariable("PATH", $Env:Path + ";C:\path\to\venv", "Machine")

  • still beeing in the jax folder in powershell actually build and compile jax

./build/build.py

  • after installation execute the command you're ask to to install jaxlib, it's something like

pip install c:\...\jax\dist\jaxlib-0.1.72-cp39-none-win_amd64.whl

  • and finally install jax with it

pip install -e .

This way it worked for me, since the documentation on the jax homepage utterly confused me.

user_169890580

Thanks to cloudhan's jax-windows-builder, it is now rather easy to install JAX and jaxlib on windows. E.g.

pip install jax==0.3.13 https://whls.blob.core.windows.net/unstable/cuda111/jaxlib-0.3.7+cuda11.cudnn82-cp38-none-win_amd64.whl

That's all.

As explained there, I had to copy the jaxlib's link address of the whl file I was interested in, i.e. the https://whls.blob.core... above).

But I also had to take care that JAX's version and Jaxlib's were compatible, which compatibility is easy to assess at github.com/google/jax/releases: just pick the last two of each (!)the version numbers! nothing to download from there.

tested with the versions explicited above. I.e. python3.8-64 & jax==0.3.13 & jaxlib-0.3.7

user_41940790

JAX does not provide jaxlib builds for Windows at this moment of time.

Issue 1, issue 2

But you can build it yourself if you wish.There are some comments in the above issue that might help you.

Admin1

This worked for me:

powershell -ExecutionPolicy ByPass -NoExit -Command "& 'C:\users\<username>\Anaconda3\condabin\conda_hook.bat' ; conda activate <yourcondaEnvironment>"
user_163464350
Related