pip installation: CPU#

Currently, the JAX team releasesjaxlib wheels for the followingoperating systems and architectures:

• Linux, x86_64

• Linux, aarch64

• macOS, Apple ARM-based

• Windows, x86_64 (experimental)

To install a CPU-only version of JAX, which might be useful for doing localdevelopment on a laptop, you can run:

pipinstall--upgradepippipinstall--upgradejax

On Windows, you may also need to install theMicrosoft Visual Studio 2019 Redistributableif it is not already installed on your machine.

Other operating systems and architectures require building from source. Tryingto pip install on other operating systems and architectures may lead tojaxlibnot being installed alongsidejax, althoughjax may successfully install(but fail at runtime).

pip installation: NVIDIA GPU (CUDA, installed via pip, easier)#

There are two ways to install JAX with NVIDIA GPU support:

• Using NVIDIA CUDA and cuDNN installed from pip wheels

• Using a self-installed CUDA/cuDNN

The JAX team strongly recommends installing CUDA and cuDNN using the pip wheels,since it is much easier!

NVIDIA has released CUDA packages only for x86_64 and aarch64.

pipinstall--upgradepip# NVIDIA CUDA 13 installation# Note: wheels only available on linux.pipinstall--upgrade"jax[cuda13]"# Alternatively, for CUDA 12, use# pip install --upgrade "jax[cuda12]"

We recommend migrating to the CUDA 13 wheels; at some point in the future wewill drop CUDA 12 support.

If JAX detects the wrong version of the NVIDIA CUDA libraries, there are several thingsyou need to check:

• Make sure thatLD_LIBRARY_PATH is not set, sinceLD_LIBRARY_PATH canoverride the NVIDIA CUDA libraries.

• Make sure that the NVIDIA CUDA libraries installed are those requested by JAX.Rerunning the installation command above should work.

pip installation: NVIDIA GPU (CUDA, installed locally, harder)#

If you prefer to use a preinstalled copy of NVIDIA CUDA, you must firstinstall NVIDIACUDA andcuDNN.

JAX provides pre-built CUDA-compatible wheels forLinux x86_64 and Linux aarch64 only. Othercombinations of operating system and architecture are possible, but requirebuilding from source (refer toBuilding from source to learn more}.

You should use an NVIDIA driver version that is at least as new as yourNVIDIA CUDA toolkit’s corresponding driver version.If you need to use a newer CUDA toolkit with an older driver, for exampleon a cluster where you cannot update the NVIDIA driver easily, you may beable to use theCUDA forward compatibility packagesthat NVIDIA provides for this purpose.

JAX currently ships two CUDA wheel variants: CUDA 12 and CUDA 13:

The CUDA 12 wheel is:

The CUDA 13 wheel is:

JAX checks the versions of your libraries, and will report an error if they arenot sufficiently new.Setting theJAX_SKIP_CUDA_CONSTRAINTS_CHECK environment variable will disablethe check, but using older versions of CUDA may lead to errors, or incorrectresults.

NCCL is an optional dependency, required only if you are performing multi-GPUcomputations.

To install, run:

pipinstall--upgradepip# Installs the wheel compatible with NVIDIA CUDA 13 and cuDNN 9.8 or newer.# Note: wheels only available on linux.pipinstall--upgrade"jax[cuda13-local]"# Installs the wheel compatible with NVIDIA CUDA 12 and cuDNN 9.8 or newer.# Note: wheels only available on linux.# pip install --upgrade "jax[cuda12-local]"

Thesepip installations do not work with Windows, and may fail silently; refer to the tableabove.

You can find your CUDA version with the command:

nvcc--version

JAX usesLD_LIBRARY_PATH to find CUDA libraries andPATH to find binaries(ptxas,nvlink). Please make sure that these paths point to the correct CUDAinstallation.

JAX requires libdevice10.bc, which typically comes from the cuda-nvvm package.Make sure that it is present in your CUDA installation.

Please let the JAX team know onthe GitHub issue trackerif you run into any errors or problems with the pre-built wheels.

NVIDIA GPU Docker containers#

NVIDIA provides theJAXToolbox containers, which arebleeding edge containers containing nightly releases of jax and somemodels/frameworks.

pip installation: Google Cloud TPU#

JAX provides pre-built wheels forGoogle Cloud TPU.To install JAX along with appropriate versions ofjaxlib andlibtpu, you can runthe following in your cloud TPU VM:

pipinstall"jax[tpu]"

For users of Colab (https://colab.research.google.com/), be sure you areusingTPU v2 and not the older, deprecated TPU runtime.

Conda installation#

There is a community-supported Conda build ofjax. To install it usingconda,simply run:

condainstalljax-cconda-forge

If you run this command on machine with an NVIDIA GPU, this should install a CUDA-enabled package ofjaxlib.

To ensure that the jax version you are installing is indeed CUDA-enabled, run:

condainstall"jaxlib=*=*cuda*"jax-cconda-forge

If you would like to override which release of CUDA is used by JAX, or toinstall the CUDA build on a machine without GPUs, follow the instructions in theTips & trickssection of theconda-forge website.

Go to theconda-forgejaxlib andjax repositoriesfor more details.