How to Install JAX with uv

Install JAX’s CPU or NVIDIA GPU packages with uv. The CUDA extras install runtime libraries from PyPI; a system CUDA Toolkit is optional for this install path.

Prerequisites

Use Python 3.12 or later in an existing uv project. Current jaxlib wheels cover CPython 3.12–3.15 on Linux x86_64/aarch64, macOS Apple Silicon, and Windows x86_64; Windows support is experimental. Intel Macs require another supported environment or a source build.

For NVIDIA GPUs, use Linux x86_64/aarch64 with a compatible NVIDIA driver. Native Windows has no CUDA wheels; WSL2 support is experimental. Consult the JAX installation matrix for AMD, Intel, or Apple GPU plugins and their separate requirements.

Install JAX for CPU

uv add jax

Expect jax and jaxlib in the installation summary. For an existing virtual environment outside a uv project, use uv pip install jax instead.

Install JAX for an NVIDIA GPU

Choose the extra by GPU architecture and driver support:

Check nvidia-smi for the driver version and GPU model. Its CUDA Version reports driver capability; it does not identify an installed toolkit.

uv add "jax[cuda13]"

Expect jax, jaxlib, jax-cuda13-plugin, jax-cuda13-pjrt, and NVIDIA runtime packages in the summary. For the CUDA 12 path, use uv add "jax[cuda12]". Neither path needs a custom package index or a system toolkit matching the extra.

If JAX loads incompatible system libraries, check LD_LIBRARY_PATH: it can override the CUDA libraries installed in the environment. Follow the JAX CUDA troubleshooting instructions before changing library paths.

Use a system CUDA Toolkit

Choose this path only when managing CUDA and cuDNN outside the Python environment:

uv add "jax[cuda13-local]"

Expect the JAX CUDA plugin and PJRT packages without NVIDIA runtime dependencies. This path requires CUDA 13+, cuDNN 9.12+ (below 10), and a compatible driver; cuda12-local requires CUDA 12.1+ and cuDNN 9.10.2+ (below 10). NCCL 2.18+ is required for multi-GPU computations.

Set PATH for the matching ptxas/nvlink binaries and LD_LIBRARY_PATH for its libraries. Include libdevice10.bc from the toolkit’s NVVM package. Check the current local-toolkit requirements before upgrading system libraries.

Install JAX for Google Cloud TPU

On a Google Cloud TPU VM:

uv add "jax[tpu]"

Expect libtpu alongside jaxlib. This does not make a CPU or GPU machine a TPU host.

Select packages by platform

For Linux CUDA hosts and macOS/Windows CPU hosts, add platform markers to the project’s dependency list:

[project]
dependencies = [
    "jax[cuda13]; sys_platform == 'linux'",
    "jax; sys_platform != 'linux'",
]

The Linux marker selects packages by operating system, not GPU presence. A GPU-less Linux image build can install and lock the CUDA dependencies; verify execution later on a compatible GPU host. Keep a separate CPU environment when CUDA dependencies are unnecessary on Linux.

Verify execution on the selected device

uv run executes the command inside the project’s virtual environment, ensuring all dependencies are installed first. For a CPU installation:

uv run python -c "import jax, jax.numpy as jnp; d=jax.devices('cpu')[0]; x=jax.device_put(jnp.arange(4), d); y=(x*x).block_until_ready(); assert int(y.sum()) == 14; print(jax.__version__, y.device, y)"

Expect a CPU device and [0 1 4 9]. For a CUDA installation, require a GPU explicitly so CPU fallback cannot pass the check:

uv run python -c "import jax, jax.numpy as jnp; d=jax.devices('gpu')[0]; x=jax.device_put(jnp.arange(4), d); y=jax.jit(lambda a: a*a)(x).block_until_ready(); assert y.device.platform == 'gpu'; assert int(y.sum()) == 14; print(jax.__version__, y.device, y)"

Expect a CUDA device and [0 1 4 9]. A missing GPU backend or failed compilation raises an error. Successful installation or import alone does not establish accelerator execution.