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.

[rgcoe]

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

[bbq572]

BEGIN: Preview Release (PR) Attempt - GPU CUDA support for python numpyro jaxlib jax - for ecommerce Bayesian Inference

this is 2025 calling, hi 2021 and 2022. Thanks to hawkinsp, I was able to complete a pure Windows CPU Build. After figuring out several env and path settings with BAZEL_LLVM being the trickiest for my free Visual Studio 2022, I was able to fully build both jaxlib, then jax latest version on 20250319. After env and path build was 100% clean. Thanks to hawkinsp, I abandoned my attempts to use MSVC compiler and used clang instead (both from free Visual Studio 2022). BAZEL_LLVM tip below.

$ pip list
Package Version Editable project location

---

absl-py 2.1.0
colorama 0.4.6
jax 0.5.3.dev20250319+1dcf872c6 C:********\jax
jaxlib 0.5.3.dev20250319
ml_dtypes 0.5.1
multipledispatch 1.0.0
numpy 2.2.4
numpyro 0.18.0
opt_einsum 3.4.0
pip 25.0.1
psutil 7.0.0
py4j 0.10.9.9
scipy 1.15.2
tqdm 4.67.1

To be safe I did the build with powershell even though I never normally use powershell for anything.

Why did I not build on Linux or WSL (Windows Subsystem For Linux). I am a Linux/Unix guy but I find that doing e-commerce with all the big firms requires me to run windows to access all their API nonsense. Their designs are so chaotic, it requires a great deal of gymnastics to get various APIs even from a single firm to beat them into submission (Amazon anyone?). Anyway, numpyro appears to be the defacto mathematician's go to library for advanced Bayesian Inference. Today I was successful in having numpyro use all available CPU core but it is clear that GPU support should be much better.

GPU support: Bottom line, I am asking for help on the best approach to tackle getting GPU CUDA working after the following successful build (hawinsp exact settings):

powershell (I included env BAZEL_LLVM to help other poor souls not have to hack as much as I had to hack):
$env:BAZEL_LLVM = "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Tools\Llvm\x64"
bazel clean --expunge
python.exe build\build.py build --wheels=jaxlib --bazel_options=--color=yes --bazel_options=--config=win_clang --verbose

QUESTION: what are best settings to attempt gpu cuda build? I expect build errors and I am willing to spend time to fix them all in order to contribute a Preview Release (PR)

Cheers!
END: Preview Release (PR) Attempt - GPU CUDA support for python numpyro jaxlib jax - for ecommerce Bayesian Inference

[hawkinsp]

We ship JAX on Windows with CPU support (only) as part of our regular releases. It is not difficult to build: https://github.com/jax-ml/jax/actions/runs/13883130724/job/38844290122

We do not currently support GPU on Windows, except via WSL2. PRs are welcome.

[hawkinsp]

(Note: we do not support MSVC as a compiler, use Clang as that Github actions job shows.)

[bbq572]

20250322 I reworded my original question above uptop to be an official offer to contribute a Preview Release (PR)

Preview Release (PR) Attempt - GPU CUDA support for python numpyro jaxlib jax - for ecommerce Bayesian Inference

Please advise/address any help as to best guesses, best approach to build as described above for GPU CUDA support

Cheers!

[bbq572]

Hi hawkinsp, in my pursuit of CUDA build,

1 of 5: IGNORE 1 of 4 as this is not an error

I could be wrong but I found a bug in the following:
The code below, if you give a path of .../bin it will return a path of .../bin/clang.exe
If you modify the code to return .../bin then this fixes the "cannot find clang.exe" error but fails elsewhere in the build
The key line of code below is
return str(pathlib.Path(which_compiler_output).resolve())
and the fix is
return str(pathlib.Path(which_compiler_output).parent.resolve())
but BOTH the original and the fixed code give different errors in the build
jax/build/tools/utils.py:

def get_compiler_path_or_exit(compiler_path_flag, compiler_name):
which_compiler_output = shutil.which(compiler_name)
if which_compiler_output:
# If we've found a compiler on the path, need to get the fully resolved path
# to ensure that system headers are found.
return str(pathlib.Path(which_compiler_output).resolve())
else:
print(
f"--{compiler_path_flag} is unset and {compiler_name} cannot be found"
f" on the PATH. Please pass --{compiler_path_flag} to the build script."
)
sys.exit(-1)

2 of 5: HARD CODED ERROR compiler = "msvc"

/external/xla/third_party/gpus/crosstool/BUILD.tpl:
cc_toolchain_config(
name = "cc-compiler-windows-config",
cpu = "x64_windows",
builtin_include_directories = [%{cxx_builtin_include_directories}],
msvc_cl_path = "%{msvc_cl_path}",
msvc_env_include = "%{msvc_env_include}",
msvc_env_lib = "%{msvc_env_lib}",
msvc_env_path = "%{msvc_env_path}",
msvc_env_tmp = "%{msvc_env_tmp}",
msvc_lib_path = "%{msvc_lib_path}",
msvc_link_path = "%{msvc_link_path}",
msvc_ml_path = "%{msvc_ml_path}",
compiler = "msvc",
)
ABOVE I AM GUESSING compiler="msvc" is incorrectly hardcoded and should be compiler="clang"

3 of 5 REPEAT CONFIGS

The following message listed multiple configs that were repeated, so I am experimenting with eliminating the repeats by commenting some of the out in build/build.py, after eliminating 3 or 4 of them the build is proceeding so I simply left the following 3 as is (repeated):
build_cuda_with_clang,
win_clang,
cuda_libraries_from_stubs

This is the build output message, building fine so far, but not done:

     "The following configs were expanded more than once: [build_cuda_with_clang, win_clang, cuda_libraries_from_stubs]. For repeatable flags, repeats are counted twice and may lead to unexpected behavior."

4 of 5 jax/jaxlib/tools:

The hack is to replace main with the C:/... absolute path for each

file: build_gpu_kernels_wheel.py
"main/jax_plugins/cuda/plugin_pyproject.toml",
"main/jax_plugins/cuda/plugin_setup.py",

"main/jaxlib/cuda/_solver.{pyext}",
guess that this may be
"C:/...../jaxlib/gpu_solver.py",

"main/jaxlib/cuda/_blas.{pyext}"
guess that this can be skipped since we guess this is merged into gpu_solver
(Basic Linear Algebra Subprograms) blas cuBLAS (nvidia)

5 of 5 #include "third_party/gpus/cuda/include/cuda.h"

not sure if edit every source to be
#include "include/cuda.h"
or if cuda.h was intended to be placed in third_party/gpus/cuda/include/
please advise as this does not build, for example,
jaxlib/cuda/cuda_plugin_extension.cc
external/xla/xla/python/py_client_gpu.h

[adam-hartshorne]

I was under the impression that not only is JAX not supported but its reliance on Triton it is now impossible to make it build for GPU support?