jax on windows

[falconair]

Hi, it looks like there has been some recent work on getting jax to work on windows. However, all the relevant tickets have been closed or merged. As far as I can tell, jax is not quite conda install jax ready for windows. Is that work being tracked somewhere?

I would love to use jax in my programming class and show students how autograd makes numpy so much more powerful! Since most of them are on windows and not yet very technically savvy, an easy windows install is a requirement.

[adam-hartshorne]

Jax itself should install via pip on Windows, it is jaxlib that needs to build from source.

I managed to get jaxlib 0.1.59 to build using the Bazel build system. The only trick I had to do was override the version of Bazel that was used. I believe Jax build files instructs a download of 3.1.0 and I had to replace it with 3.7.0 (and rename as if it was 3.1.0) in /build directory.

[ryancoe]

I've been using autograd within a project I'm working on, but since that is now unsupported I would like to transition to jax if possible without requiring Windows users to compile from source or use an unsupported binary. I understand that the jax + jaxlib on Windows currently requires users to compile from source, but I'm a bit unclear if there's a clean and reliable way to install just the CPU parts of jax on Windows. All I need is the jax.numpy module.

I am able to do pip install jax[cpu] on Windows, but I get the following error when calling import jax.numpy. Is there a good solution for this that I'm missing?

In [1]: import jax.numpy as jnp
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
~\Anaconda3\envs\jax\lib\site-packages\jax\_src\lib\__init__.py in <module>
     37 try:
---> 38   import jaxlib as jaxlib
     39 except ModuleNotFoundError as err:

ModuleNotFoundError: No module named 'jaxlib'

The above exception was the direct cause of the following exception:

ModuleNotFoundError                       Traceback (most recent call last)
<ipython-input-1-cb946e78fdef> in <module>
----> 1 import jax.numpy as jnp

~\Anaconda3\envs\jax\lib\site-packages\jax\__init__.py in <module>
     35 # We want the exported object to be the class, so we first import the module
     36 # to make sure a later import doesn't overwrite the class.
---> 37 from jax import config as _config_module
     38 del _config_module
     39

~\Anaconda3\envs\jax\lib\site-packages\jax\config.py in <module>
     16
     17 # flake8: noqa: F401
---> 18 from jax._src.config import config

~\Anaconda3\envs\jax\lib\site-packages\jax\_src\config.py in <module>
     25 import warnings
     26
---> 27 from jax._src import lib
     28 from jax._src.lib import jax_jit
     29 if lib.xla_extension_version >= 58:

~\Anaconda3\envs\jax\lib\site-packages\jax\_src\lib\__init__.py in <module>
     38   import jaxlib as jaxlib
     39 except ModuleNotFoundError as err:
---> 40   raise ModuleNotFoundError(
     41     'jax requires jaxlib to be installed. See '
     42     'https://github.com/google/jax#installation for installation instructions.'

ModuleNotFoundError: jax requires jaxlib to be installed. See https://github.com/google/jax#installation for installation instructions.

[adam-hartshorne]

There are now pre-built versions of JAX for windows available here https://github.com/cloudhan/jax-windows-builder

[falconair]

Looks like there is a ticket exactly relevant to this topic: #5795