winjax
Native Windows CUDA support for JAX — no WSL2 required. Unofficial.
winjax installs a PJRT GPU plugin built
natively for Windows from XLA sources, plus a small loader that registers it
with stock jax/jaxlib. All CUDA runtime libraries (CUDA 13, cuDNN 9) come
from NVIDIA's pip wheels, so a working NVIDIA driver is the only system
requirement.
Requirements
- Windows 10/11, x86-64
- Python 3.13
- An NVIDIA GPU with a driver supporting CUDA 13
Install
pip install winjax
Use
import jax
print(jax.devices()) # [CudaDevice(id=0)]
Nothing else to configure: import jax discovers the plugin through the
jax_plugins namespace package.
Packages
winjax— loader (this package, pure Python)winjax-cuda13-pjrt— the Windows-built XLA CUDA PJRT plugin DLLwinjax-cuda13-plugin— CUDA kernel extension modules (jax_cuda13_plugin)
Source and patches: https://github.com/eterevsky/winjax
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file winjax-0.11.0-py3-none-any.whl.
File metadata
- Download URL: winjax-0.11.0-py3-none-any.whl
- Upload date:
- Size: 5.1 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: uv/0.11.8 {"installer":{"name":"uv","version":"0.11.8","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":null,"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
7d5b7659316a7be6f05a1395098d6f163aeff77ce32cba823f990f3426447341
|
|
| MD5 |
6060a03518c9c6a3374034f55bbf16bf
|
|
| BLAKE2b-256 |
b6a3d019a89df09a77d7fe4b2565651e26123ab3de4392000f27d11140f146c1
|
