HajimeKawahara / exojax

🐈 Automatic differentiable spectrum modeling of exoplanets/brown dwarfs using JAX, compatible with NumPyro and JAXopt
http://secondearths.sakura.ne.jp/exojax/
MIT License
57 stars 14 forks source link

JAX does not install on Windows #198

Closed dcmvdbekerom closed 2 years ago

dcmvdbekerom commented 2 years ago

Hello,

Upon trying to install exojax (pip install exojax), pip is installing a large number of different versions. ~It appears as though this is because of conflicting dependency requirements. Is this a problem of my system or of exojax? This behaviour appears both in python 3.7 and anaconda (python 3.8), both on Windows 10.~

EDIT: The issue is that jax, or more specifically jaxlib, is not currently available for Windows. There are instructions online how to build it yourself, I'll report back here when I'm succesful.

One could also try prebuilt version here: https://github.com/cloudhan/jax-windows-builder but I haven't tried this yet.

(base) C:\Users\dcmvd>pip install exojax
Collecting exojax
  Using cached exojax-1.0.0-py2.py3-none-any.whl (450 kB)
Requirement already satisfied: tqdm in c:\users\dcmvd\anaconda3\lib\site-packages (from exojax) (4.59.0)
Requirement already satisfied: numpy in c:\users\dcmvd\anaconda3\lib\site-packages (from exojax) (1.20.1)
Requirement already satisfied: pyarrow in c:\users\dcmvd\anaconda3\lib\site-packages (from exojax) (6.0.0)
Collecting jax
  Using cached jax-0.2.26.tar.gz (850 kB)
Requirement already satisfied: scipy in c:\users\dcmvd\anaconda3\lib\site-packages (from exojax) (1.6.2)
Collecting bs4
  Using cached bs4-0.0.1.tar.gz (1.1 kB)
Requirement already satisfied: pandas in c:\users\dcmvd\anaconda3\lib\site-packages (from exojax) (1.2.4)
Collecting numpyro
  Using cached numpyro-0.8.0-py3-none-any.whl (264 kB)
Requirement already satisfied: vaex in c:\users\dcmvd\anaconda3\lib\site-packages (from exojax) (4.5.0)
Requirement already satisfied: beautifulsoup4 in c:\users\dcmvd\anaconda3\lib\site-packages (from bs4->exojax) (4.9.3)
Requirement already satisfied: soupsieve>1.2 in c:\users\dcmvd\anaconda3\lib\site-packages (from beautifulsoup4->bs4->exojax) (2.2.1)
Collecting absl-py
  Using cached absl_py-1.0.0-py3-none-any.whl (126 kB)
Collecting opt_einsum
  Using cached opt_einsum-3.3.0-py3-none-any.whl (65 kB)
Requirement already satisfied: typing_extensions in c:\users\dcmvd\anaconda3\lib\site-packages (from jax->exojax) (3.10.0.2)
Requirement already satisfied: six in c:\users\dcmvd\anaconda3\lib\site-packages (from absl-py->jax->exojax) (1.15.0)
Collecting numpyro
  Using cached numpyro-0.7.2-py3-none-any.whl (250 kB)
  Using cached numpyro-0.7.1-py3-none-any.whl (250 kB)
  Using cached numpyro-0.7.0-py3-none-any.whl (247 kB)
  Using cached numpyro-0.6.0-py3-none-any.whl (218 kB)
Collecting jax
  Using cached jax-0.2.10.tar.gz (589 kB)
Collecting numpyro
  Using cached numpyro-0.5.0-py3-none-any.whl (194 kB)
Collecting jax
  Using cached jax-0.2.8.tar.gz (544 kB)
Collecting numpyro
  Using cached numpyro-0.4.1-py3-none-any.whl (176 kB)
  Using cached numpyro-0.4.0-py3-none-any.whl (175 kB)
Collecting jax
  Using cached jax-0.2.0.tar.gz (454 kB)
Collecting numpyro
  Using cached numpyro-0.3.0-py3-none-any.whl (157 kB)
Collecting jax
  Using cached jax-0.1.73.tar.gz (419 kB)
Collecting numpyro
  Using cached numpyro-0.2.4-py3-none-any.whl (159 kB)
Collecting jax
  Using cached jax-0.1.57.tar.gz (255 kB)
Collecting numpyro
  Using cached numpyro-0.2.3-py3-none-any.whl (145 kB)
  Using cached numpyro-0.2.2-py3-none-any.whl (145 kB)
  Using cached numpyro-0.2.1-py3-none-any.whl (141 kB)
  Using cached numpyro-0.2.0-py3-none-any.whl (95 kB)
  Using cached numpyro-0.1.0.tar.gz (73 kB)
Collecting jax
  Using cached jax-0.1.35.tar.gz (181 kB)
INFO: pip is looking at multiple versions of absl-py to determine which version is compatible with other requirements. This could take a while.
Collecting absl-py
  Using cached absl_py-0.15.0-py3-none-any.whl (132 kB)
  Using cached absl_py-0.14.1-py3-none-any.whl (131 kB)
  Using cached absl_py-0.14.0-py3-none-any.whl (131 kB)
  Using cached absl_py-0.13.0-py3-none-any.whl (132 kB)
  Using cached absl_py-0.12.0-py3-none-any.whl (129 kB)
  Using cached absl_py-0.11.0-py3-none-any.whl (127 kB)
  Using cached absl_py-0.10.0-py3-none-any.whl (127 kB)
  Using cached absl-py-0.9.0.tar.gz (104 kB)
  Using cached absl-py-0.8.1.tar.gz (103 kB)
  Using cached absl-py-0.8.0.tar.gz (102 kB)
  Using cached absl-py-0.7.1.tar.gz (99 kB)
  Using cached absl-py-0.7.0.tar.gz (96 kB)
  Using cached absl-py-0.6.1.tar.gz (94 kB)
  Using cached absl-py-0.6.0.tar.gz (93 kB)
  Using cached absl-py-0.5.0.tar.gz (90 kB)
  Using cached absl-py-0.4.1.tar.gz (88 kB)
  Using cached absl-py-0.4.0.tar.gz (88 kB)
  Using cached absl-py-0.3.0.tar.gz (84 kB)
  Using cached absl-py-0.2.2.tar.gz (82 kB)
  Using cached absl-py-0.2.1.tar.gz (81 kB)
  Using cached absl-py-0.2.0.tar.gz (82 kB)
  Using cached absl-py-0.1.13.tar.gz (80 kB)
  Using cached absl-py-0.1.12.tar.gz (79 kB)
  Using cached absl-py-0.1.11.tar.gz (80 kB)
  Using cached absl-py-0.1.10.tar.gz (79 kB)
  Using cached absl-py-0.1.9.tar.gz (79 kB)
  Using cached absl-py-0.1.8.tar.gz (79 kB)
  Using cached absl-py-0.1.7.tar.gz (78 kB)
  Using cached absl-py-0.1.6.tar.gz (78 kB)
  Using cached absl-py-0.1.5.tar.gz (78 kB)
  Using cached absl-py-0.1.4.tar.gz (78 kB)
  Using cached absl-py-0.1.3.tar.gz (78 kB)
  Using cached absl-py-0.1.2.tar.gz (78 kB)
  Using cached absl-py-0.1.1.tar.gz (77 kB)
  Using cached absl-py-0.1.0.tar.gz (76 kB)
INFO: pip is looking at multiple versions of absl-py to determine which version is compatible with other requirements. This could take a while.
INFO: This is taking longer than usual. You might need to provide the dependency resolver with stricter constraints to reduce runtime. If you want to abort this run, you can press Ctrl + C to do so. To improve how pip performs, tell us what happened here: https://pip.pypa.io/surveys/backtracking
INFO: pip is looking at multiple versions of scipy to determine which version is compatible with other requirements. This could take a while.
Collecting scipy
  Using cached scipy-1.7.3-cp38-cp38-win_amd64.whl (34.2 MB)
  Using cached scipy-1.7.2-cp38-cp38-win_amd64.whl (34.2 MB)
  Using cached scipy-1.7.1-cp38-cp38-win_amd64.whl (33.7 MB)
  Downloading scipy-1.7.0-cp38-cp38-win_amd64.whl (33.7 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 33.7 MB 6.4 MB/s
  Downloading scipy-1.6.3-cp38-cp38-win_amd64.whl (32.7 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 32.7 MB 6.4 MB/s
  Downloading scipy-1.6.1-cp38-cp38-win_amd64.whl (32.7 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 32.7 MB 6.4 MB/s
  Downloading scipy-1.6.0-cp38-cp38-win_amd64.whl (32.7 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 32.7 MB 3.2 MB/s
  Downloading scipy-1.5.4-cp38-cp38-win_amd64.whl (31.4 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 31.4 MB 3.3 MB/s
  Downloading scipy-1.5.3-cp38-cp38-win_amd64.whl (31.4 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 31.4 MB 6.8 MB/s
  Downloading scipy-1.5.2-cp38-cp38-win_amd64.whl (31.4 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 31.4 MB 6.8 MB/s
  Downloading scipy-1.5.1-cp38-cp38-win_amd64.whl (31.4 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 31.4 MB 3.3 MB/s
  Downloading scipy-1.5.0-cp38-cp38-win_amd64.whl (31.4 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 31.4 MB 3.3 MB/s
  Downloading scipy-1.4.1-cp38-cp38-win_amd64.whl (31.0 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 31.0 MB 3.2 MB/s
  Downloading scipy-1.4.0-cp38-cp38-win_amd64.whl (31.0 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 31.0 MB 312 kB/s
  Downloading scipy-1.3.3-cp38-cp38-win_amd64.whl (30.7 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 30.7 MB 3.3 MB/s
  Downloading scipy-1.3.2-cp38-cp38-win_amd64.whl (30.7 MB)
     |β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 30.7 MB 3.3 MB/s
erwanp commented 2 years ago

Indeed, no jax on Windows. I eventually installed Jax on a Linux from my virtualmachine, but from there I couldn't properly access the GPU. You may want to give a try to Google Colab which has some GPU hardware available

dcmvdbekerom commented 2 years ago

I tried building jaxlib myself on Windows to no avail. Ultimately I was able to install it using jaxlib prebuilt for windows, found here: https://whls.blob.core.windows.net/unstable/index.html This worked, I now have jax, numpyro, and exojax working!