jax windows compatibility?

[omarcr]

In the readme there is not a single mention of compatibility with windows. Is it possible? I am trying to install the library but I get the following error when installing:

C:\Users\Omar\Downloads\jax-master>python build/build.py --enable_cuda

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Bazel binary path: C:\path_apps\bazel.EXE
Python binary path: C:\Program Files\Python36\python.exe
MKL-DNN enabled: yes
-march=native: no
CUDA enabled: yes
CUDA toolkit path: /usr/local/cuda
CUDNN library path: /usr/local/cuda

Building XLA and installing it in the jaxlib source tree...
ERROR: python_path must be an absolute path when it is set.
INFO: Elapsed time: 0.123s
INFO: 0 processes.
FAILED: Build did NOT complete successfully (0 packages loaded)
FAILED: Build did NOT complete successfully (0 packages loaded)
Traceback (most recent call last):
  File "build/build.py", line 317, in <module>
  File "build/build.py", line 313, in main
  File "build/build.py", line 50, in shell
  File "C:\Program Files\Python36\lib\subprocess.py", line 356, in check_output
    **kwargs).stdout
  File "C:\Program Files\Python36\lib\subprocess.py", line 438, in run
    output=stdout, stderr=stderr)
subprocess.CalledProcessError: Command '['C:\\path_apps\\bazel.EXE', 'run', '--verbose_failures=true', '--config=mkl_open_source_only', '--config=cuda', ':install_xla_in_source_tree', 'C:\\Users\\Omar\\Downloads\\jax-master\\build']' returned non-zero exit status 2.

[hawkinsp]

Jax doesn't support windows yet. We don't have the engineering time to work on it ourselves, but contributions are welcome!

It should be quite feasible to do — most of the C++ parts of Jax (Jaxlib) are shared with TensorFlow and I would imagine build on Windows already, if the build scripts were to be configured correctly.

Closing as this is a duplicate of #438.

[omarcr]

please state on read me that is not compatible with windows since tensorflow, CUDA, CUDNN etc are all now compatible so I assumed this is as well for JAX but is not the case (yet). I hope you can make this compatible with windows soon. Thanks.

[mattjj]

Good idea; done in 3a633db.