Installingjaxlib with pip#

If you’re only modifying Python portions of JAX, we recommend installingjaxlib from a prebuilt wheel using pip:

pipinstalljaxlib

See theJAX readme for fullguidance on pip installation (e.g., for GPU and TPU support).

Buildingjaxlib from source#

To buildjaxlib from source, you must also install some prerequisites:

• A C++ compiler:

  As mentioned in the box above, it is best to use a recent version of clang(at the time of writing, the version we test is 18), but other compilers (e.g.g++ or MSVC) may work.

  On Ubuntu or Debian you can follow the instructions from theLLVM documentation to install the latest stableversion of clang.

  If you are building on a Mac, make sure XCode and the XCode command line toolsare installed.

  See below for Windows build instructions.

• Python: for running the build helper script. Note that there is no need toinstall Python dependencies locally, as your system Python will be ignoredduring the build; please checkManaging hermetic Python for details.

To buildjaxlib for CPU or TPU, you can run:

pythonbuild/build.pybuild--wheels=jaxlib--verbosepipinstalldist/*.whl# installs jaxlib (includes XLA)

To build a wheel for a version of Python different from your current systeminstallation pass--python_version flag to the build command:

pythonbuild/build.pybuild--wheels=jaxlib--python_version=3.12--verbose

The rest of this document assumes that you are building for Python versionmatching your current system installation. If you need to build for a differentversion, simply append--python_version=<pyversion> flag every time you callpythonbuild/build.py. Note, the Bazel build will always use a hermetic Pythoninstallation regardless of whether the--python_version parameter is passed ornot.

If you would like to buildjaxlib and the CUDA plugins: Run

pythonbuild/build.pybuild--wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt

to generate three wheels (jaxlib without cuda, jax-cuda-plugin, andjax-cuda-pjrt). By default all CUDA compilation steps performed by NVCC andclang, but it can be restricted to clang via the--build_cuda_with_clang flag.

Seepythonbuild/build.py--help for configuration options. Herepython should be the name of your Python 3 interpreter; on some systems, youmay need to usepython3 instead. Despite calling the script withpython,Bazel will always use its own hermetic Python interpreter and dependencies, onlythebuild/build.py script itself will be processed by your system Pythoninterpreter. By default, the wheel is written to thedist/ subdirectory of thecurrent directory.

• JAX versions starting from v.0.4.32: you can provide custom CUDA and CUDNNversions in the configuration options. Bazel will download them and use astarget dependencies.

  To download the specific versions of CUDA/CUDNN redistributions, you can usethe--cuda_version and--cudnn_version flags:

  pythonbuild/build.pybuild--wheels=jax-cuda-plugin--cuda_version=12.3.2\--cudnn_version=9.1.1

  or

  pythonbuild/build.pybuild--wheels=jax-cuda-pjrt--cuda_version=12.3.2\--cudnn_version=9.1.1

  Please note that these parameters are optional: by default Bazel willdownload CUDA and CUDNN redistribution versions provided in.bazelrc in theenvironment variablesHERMETIC_CUDA_VERSION andHERMETIC_CUDNN_VERSIONrespectively.

  To point to CUDA/CUDNN/NCCL redistributions on local file system, you can usethe following command:

  pythonbuild/build.pybuild--wheels=jax-cuda-plugin\--bazel_options=--repo_env=LOCAL_CUDA_PATH="/foo/bar/nvidia/cuda"\--bazel_options=--repo_env=LOCAL_CUDNN_PATH="/foo/bar/nvidia/cudnn"\--bazel_options=--repo_env=LOCAL_NCCL_PATH="/foo/bar/nvidia/nccl"

  Please see the full list of instructions inXLA documentation.

• JAX versions prior v.0.4.32: you must have CUDA and CUDNN installed andprovide paths to them using configuration options.

Building jaxlib from source with a modified XLA repository.#

JAX depends on XLA, whose source code is in theXLA GitHub repository.By default JAX uses a pinned copy of the XLA repository, but we oftenwant to use a locally-modified copy of XLA when working on JAX. There are twoways to do this:

• use Bazel’soverride_repository feature, which you can pass as a commandline flag tobuild.py as follows:

  pythonbuild/build.pybuild--wheels=jaxlib--local_xla_path=/path/to/xla

• modify theWORKSPACE file in the root of the JAX source tree to point toa different XLA tree.

To contribute changes back to XLA, send PRs to the XLA repository.

The version of XLA pinned by JAX is regularly updated, but is updated inparticular before eachjaxlib release.

Additional Notes for Buildingjaxlib from source on Windows#

Note: JAX does not support CUDA on Windows; use WSL2 for CUDA support.

On Windows, followInstall Visual Studioto set up a C++ toolchain. Visual Studio 2019 version 16.5 or newer is required.

JAX builds use symbolic links, which require that you activateDeveloper Mode.

You can either install Python using itsWindows installer, or if you prefer, youcan useAnacondaorMinicondato set up a Python environment.

Some targets of Bazel use bash utilities to do scripting, soMSYS2is needed. SeeInstalling Bazel on Windowsfor more details. Install the following packages:

pacman-Spatchcoreutils

Once coreutils is installed, the realpath command should be present in your shell’s path.

Once everything is installed. Open PowerShell, and make sure MSYS2 is in thepath of the current session. Ensurebazel,patch andrealpath areaccessible. Activate the conda environment.

python.\build\build.pybuild--wheels=jaxlib

To build with debug information, add the flag--bazel_options='--copt=/Z7'.

Additional notes for building a ROCMjaxlib for AMD GPUs#

For detailed instructions on buildingjaxlib with ROCm support, refer to the official guide:Build ROCm JAX from Source

Specifying Python version#

When you runbuild/build.py tool, the version of hermetic Python is setautomatically to match the version of the Python you used torunbuild/build.py script. To choose a specific version explicitly you maypass--python_version argument to the tool:

pythonbuild/build.pybuild--python_version=3.12

Under the hood, the hermetic Python version is controlledbyHERMETIC_PYTHON_VERSION environment variable, which is set automaticallywhen you runbuild/build.py. In case you run bazel directly you may need toset the variable explicitly in one of the following ways:

# Either add an entry to your `.bazelrc` filebuild--repo_env=HERMETIC_PYTHON_VERSION=3.12# OR pass it directly to your specific build commandbazelbuild<target>--repo_env=HERMETIC_PYTHON_VERSION=3.12# OR set the environment variable globally in your shell:exportHERMETIC_PYTHON_VERSION=3.12

You may run builds and tests against different versions of Python sequentiallyon the same machine by simply switching the value of--python_version betweenthe runs. All the python-agnostic parts of the build cache from the previousbuild will be preserved and reused for the subsequent builds.

Specifying Python dependencies#

During bazel build all JAX’s Python dependencies are pinned to their specificversions. This is necessary to ensure reproducibility of the build.The pinned versions of the full transitive closure of JAX’s dependenciestogether with their corresponding hashes are specified inbuild/requirements_lock_<pythonversion>.txt files (e.g.build/requirements_lock_3_12.txt forPython3.12).

To update the lock files, make surebuild/requirements.in contains the desireddirect dependencies list and then execute the following command (which will callpip-compile under the hood):

pythonbuild/build.pyrequirements_update--python_version=3.12

Alternatively, if you need more control, you may run the bazel commanddirectly (the two commands are equivalent):

# Regular Pythonbazelrun//build:requirements.update--repo_env=HERMETIC_PYTHON_VERSION=3.12# Free-threaded Pythonbazelrun//build:requirements_ft.update--repo_env=HERMETIC_PYTHON_VERSION=3.13-ft

Note, since it is stillpip andpip-compile tools used under the hood, somost of the command line arguments and features supported by those tools will beacknowledged by the Bazel requirements updater command as well. For example, ifyou wish the updater to consider pre-release versions simply pass--preargument to the bazel command:

bazelrun//build:requirements.update--repo_env=HERMETIC_PYTHON_VERSION=3.12----pre

Specifying dependencies on local wheels#

By default the build scansdist directory in the repository root for any local.whl files to be included in the list of dependencies. If the wheel is Pythonversion specific, only the wheels that match the selected Python version willbe included.

The overall local wheel search and selection logic is controlled by thearguments topython_init_repositories() macro (called directly from theWORKSPACE file). You may uselocal_wheel_dist_folder to change the locationof the folder with local wheels. Uselocal_wheel_inclusion_list andlocal_wheel_exclusion_list arguments to specify which wheels should beincluded and/or excluded from the search (it supports basic wildcard matching).

If necessary, you can also depend on a local.whl file manually, bypassing theautomatic local wheel search mechanism. For example to depend on your newlybuilt jaxlib wheel, you may add a path to the wheel inbuild/requirements.inand re-run the requirements updater command for a selected version of Python.For example:

echo-e"\n$(realpath jaxlib-0.4.27.dev20240416-cp312-cp312-manylinux_2_27_x86_64.whl)">>build/requirements.inpythonbuild/build.pyrequirements_update--python_version=3.12

Specifying dependencies on nightly wheels#

To build and test against the very latest, potentially unstable, set of Pythondependencies we provide a special version of the dependency updater command asfollows:

pythonbuild/build.pyrequirements_update--python_version=3.12--nightly_update

Or, if you runbazel directly (the two commands are equivalent):

bazelrun//build:requirements_nightly.update--repo_env=HERMETIC_PYTHON_VERSION=3.12

The difference between this and the regular updater is that by default it wouldaccept pre-release, dev and nightly packages, it will alsosearch https://pypi.anaconda.org/scientific-python-nightly-wheels/simple as anextra index url and will not put hashes in the resultant requirements lock file.

Customizing hermetic Python (Advanced Usage)#

We support all of the current versions of Python out of the box, so unless yourworkflow has very special requirements (such as ability to use your own customPython interpreter) you may safely skip this section entirely.

In short, if you rely on a non-standard Python workflow you still can achievethe great level of flexibility in hermetic Python setup. Conceptually there willbe only one difference compared to non-hermetic case: you will need to think interms of files, not installations (i.e. think what files your build actuallydepends on, not what files need to be installed on your system), the rest ispretty much the same.

So, in practice, to gain full control over your Python environment, hermetic ornot you need to be able to do the following three things:

1. Specify which python interpreter to use (i.e. pick actualpython orpython3 binary and libs that come with it in the same folder).

2. Specify a list of Python dependencies (e.g.numpy) and their actualversions.

3. Be able to add/remove/update dependencies in the list easily. Eachdependency itself could be custom too (self-built for example).

You already know how to do all of the steps above in a non-hermetic Pythonenvironment, here is how you do the same in the hermetic one (by approaching itin terms of files, not installations):

1. Instead of installing Python, get Python interpreter in atar orzipfile. Depending on your case you may simply pull one of many existing ones(such aspython-build-standalone),or build your own and pack it in an archive (following officialbuild instructionswill do just fine). E.g. on Linux it will look something like the following:

   ./configure--prefixpythonmake-j12makealtinstalltar-czpfmy_python.tgzpython

   Once you have the tarball ready, plug it in the build by pointingHERMETIC_PYTHON_URL env var to the archive (either local one or from theinternet):

   --repo_env=HERMETIC_PYTHON_URL="file:///local/path/to/my_python.tgz"--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum># OR--repo_env=HERMETIC_PYTHON_URL="https://remote/url/to/my_python.tgz"--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum># We assume that top-level folder in the tarball is called "python", if it is# something different just pass additional HERMETIC_PYTHON_PREFIX parameter--repo_env=HERMETIC_PYTHON_URL="https://remote/url/to/my_python.tgz"--repo_env=HERMETIC_PYTHON_SHA256=<file's_sha256_sum>--repo_env=HERMETIC_PYTHON_PREFIX="my_python/install"

2. Instead of doingpipinstall createrequirements_lock.txt file withfull transitive closure of your dependencies. You may also depend on theexisting ones already checked in this repo (as long as they work with yourcustom Python version). There are no special instructions on how you do it,you may follow steps recommended inSpecifying Python dependenciesfrom this doc, just call pip-compile directly (note, the lock file must behermetic, but you can always generate it from non-hermetic python if you’dlike) or even create it manually (note, hashes are optional in lock files).

3. If you need to update or customize your dependencies list, you may once againfollow theSpecifying Python dependenciesinstructions to updaterequirements_lock.txt, call pip-compile directly ormodify it manually. If you have a custom package you want to use just pointto its.whl file directly (remember, work in terms of files, notinstallations) from your lock (note,requirements.txt andrequirements_lock.txt files support local wheel references). If yourrequirements_lock.txt is already specified as a dependency topython_init_repositories() inWORKSPACE file you don’t have to doanything else. Otherwise you can point to your custom file as follows:

   --repo_env=HERMETIC_REQUIREMENTS_LOCK="/absolute/path/to/custom_requirements_lock.txt"

   Also note if you useHERMETIC_REQUIREMENTS_LOCK then it fully controls listof your dependencies and the automatic local wheels resolution logicdescribed inSpecifying dependencies on local wheelsgets disabled to not interfere with it.

That is it. To summarize: if you have an archive with Python interpreter in itand a requirements_lock.txt file with full transitive closure of yourdependencies then you fully control your Python environment.

bazelbuild<target>--repo_env=HERMETIC_PYTHON_VERSION=3.13--repo_env=HERMETIC_PYTHON_URL="https://github.com/indygreg/python-build-standalone/releases/download/20241016/cpython-3.13.0+20241016-x86_64-unknown-linux-gnu-install_only.tar.gz"--repo_env=HERMETIC_PYTHON_SHA256="2c8cb15c6a2caadaa98af51df6fe78a8155b8471cb3dd7b9836038e0d3657fb4"

bazeltest<target>--repo_env=HERMETIC_PYTHON_VERSION=3.13--repo_env=HERMETIC_PYTHON_URL="file:///path/to/cpython.tar.gz"--repo_env=HERMETIC_PYTHON_PREFIX="prefix/to/strip/in/cython/tar/gz/archive"--repo_env=HERMETIC_PYTHON_SHA256=<sha256_sum>--repo_env=HERMETIC_REQUIREMENTS_LOCK="/absolute/path/to/build:custom_requirements_lock.txt"

bazeltest<target>--repo_env=HERMETIC_PYTHON_VERSION=3.13--repo_env=HERMETIC_REQUIREMENTS_LOCK="/absolute/path/to/build:custom_requirements_lock.txt"

requirements={"3.11":"//build:requirements_lock_3_11.txt","3.12":"//build:requirements_lock_3_12.txt","3.13":"//build:requirements_lock_3_13.txt","3.13-scenario1":"//build:scenario1_requirements_lock_3_13.txt","3.13-scenario2":"//build:scenario2_requirements_lock_3_13.txt",},

# To build with scenario1 dependencies:bazeltest<target>--repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario1# To build with scenario2 dependencies:bazeltest<target>--repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario2# To build with default dependencies:bazeltest<target>--repo_env=HERMETIC_PYTHON_VERSION=3.13# To build with scenario1 dependencies and custom Python 3.13 interpreter:bazeltest<target>--repo_env=HERMETIC_PYTHON_VERSION=3.13-scenario1--repo_env=HERMETIC_PYTHON_URL="file:///path/to/cpython.tar.gz"--repo_env=HERMETIC_PYTHON_SHA256=<sha256_sum>

Using Bazel#

First, configure the JAX build by using the--configure_only flag. Pass--wheel_list=jaxlib for CPU tests and CUDA/ROCM for GPU for GPU tests:

pythonbuild/build.pybuild--wheels=jaxlib--configure_onlypythonbuild/build.pybuild--wheels=jax-cuda-plugin--configure_onlypythonbuild/build.pybuild--wheels=jax-rocm-plugin--configure_only

You may pass additional options tobuild.py to configure the build; see thejaxlib build documentation for details.

By default the Bazel build runs the JAX tests usingjaxlib built from source.To run JAX tests, run:

bazeltest//tests:cpu_tests//tests:backend_independent_tests

//tests:gpu_tests and//tests:tpu_tests are also available, if you have thenecessary hardware.

You need to configurecuda to rungpu tests:

pythonbuild/build.pybuild--wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt--configure_only

To use a preinstalledjaxlib instead of building it you first need tomake it available in the hermetic Python. To install a specific version ofjaxlib within hermetic Python run (usingjaxlib>=0.4.26 as an example):

echo-e"\njaxlib >= 0.4.26">>build/requirements.inpythonbuild/build.pyrequirements_update

Alternatively, to installjaxlib from a local wheel (assuming Python 3.12):

echo-e"\n$(realpath jaxlib-0.4.26-cp312-cp312-manylinux_2_27_x86_64.whl)">>build/requirements.inpythonbuild/build.pyrequirements_update--python_version=3.12

Once you havejaxlib installed hermetically, run:

bazeltest--//jax:build_jaxlib=false//tests:cpu_tests//tests:backend_independent_tests

A number of test behaviors can be controlled using environment variables (seebelow). Environment variables may be passed to JAX tests using the--test_env=FLAG=value flag to Bazel.

Some of JAX tests are for multiple accelerators (i.e. GPUs, TPUs). When JAX isalready installed, you can run GPUs tests like this:

bazeltest//tests:gpu_tests--local_test_jobs=4--test_tag_filters=multiaccelerator--//jax:build_jaxlib=false--test_env=XLA_PYTHON_CLIENT_ALLOCATOR=platform

You can speed up single accelerator tests by running them in parallel onmultiple accelerators. This also triggers multiple concurrent tests peraccelerator. For GPUs, you can do it like this:

NB_GPUS=2JOBS_PER_ACC=4J=$((NB_GPUS * JOBS_PER_ACC))MULTI_GPU="--run_under $PWD/build/parallel_accelerator_execute.sh --test_env=JAX_ACCELERATOR_COUNT=${NB_GPUS} --test_env=JAX_TESTS_PER_ACCELERATOR=${JOBS_PER_ACC} --local_test_jobs=$J"bazel test //tests:gpu_tests //tests:backend_independent_tests --test_env=XLA_PYTHON_CLIENT_PREALLOCATE=false --test_tag_filters=-multiaccelerator $MULTI_GPU

Usingpytest#

First, install the dependencies byrunningpipinstall-rbuild/test-requirements.txt.

To run all the JAX tests usingpytest, we recommend usingpytest-xdist,which can run tests in parallel. It is installed as a part ofpipinstall-rbuild/test-requirements.txt command.

From the repository root directory run:

pytest-nautotests

Controlling test behavior#

JAX generates test cases combinatorially, and you can control the number ofcases that are generated and checked for each test (default is 10) using theJAX_NUM_GENERATED_CASES environment variable. The automated testscurrently use 25 by default.

For example, one might write

# Bazelbazel test //tests/... --test_env=JAX_NUM_GENERATED_CASES=25`

or

# pytestJAX_NUM_GENERATED_CASES=25pytest-nautotests

The automated tests also run the tests with default 64-bit floats and ints(JAX_ENABLE_X64):

JAX_ENABLE_X64=1JAX_NUM_GENERATED_CASES=25pytest-nautotests

You can run a more specific set of tests usingpytest’sbuilt-in selection mechanisms, or alternatively you can run a specific testfile directly to see more detailed information about the cases being run:

JAX_NUM_GENERATED_CASES=5pythontests/lax_numpy_test.py

You can skip a few tests known to be slow, by passing environment variableJAX_SKIP_SLOW_TESTS=1.

To specify a particular set of tests to run from a test file, you can pass a stringor regular expression via the--test_targets flag. For example, you can run allthe tests ofjax.numpy.pad using:

pythontests/lax_numpy_test.py--test_targets="testPad"

The Colab notebooks are tested for errors as part of the documentation build.

Hypothesis tests#

Some of the tests usehypothesis.Normally, hypothesis will test using multiple example inputs, and on a test failureit will try to find a smaller example that still results in failure:Look through the test failure for a line like the one below, and add the decoratormentioned in the message:

Youcanreproducethisexamplebytemporarilyadding@reproduce_failure('6.97.4',b'AXicY2DAAAAAEwAB')asadecoratoronyourtestcase

For interactive development, you can set the environment variableJAX_HYPOTHESIS_PROFILE=interactive (or the equivalent flag--jax_hypothesis_profile=interactive)in order to set the number of examples to 1, and skip the exampleminimization phase.

Doctests#

JAX uses pytest in doctest mode to test the code examples within the documentation.You can find the up-to-date command to run doctests inci-build.yaml.E.g., you can run:

JAX_TRACEBACK_FILTERING=offXLA_FLAGS=--xla_force_host_platform_device_count=8pytest-nauto--tb=short--doctest-glob='*.md'--doctest-glob='*.rst'docs--doctest-continue-on-failure--ignore=docs/multi_process.md

Additionally, JAX runs pytest indoctest-modules mode to ensure code examples infunction docstrings will run correctly. You can run this locally using, for example:

JAX_TRACEBACK_FILTERING=offXLA_FLAGS=--xla_force_host_platform_device_count=8pytest--doctest-modulesjax/_src/numpy/lax_numpy.py

Update notebooks#

We usejupytext to maintain two synced copies of the notebooksindocs/notebooks: one inipynb format, and one inmd format. The advantage of the formeris that it can be opened and executed directly in Colab; the advantage of the latter is thatit makes it much easier to track diffs within version control.

pipinstalljupytext==1.16.4jupytext--syncdocs/notebooks/thinking_in_jax.ipynb

pipinstallpre-commitpre-commitrunjupytext--all-files

jupytext--set-formatsipynb,md:mystpath/to/the/notebook.ipynb

Documentation building onreadthedocs.io#

JAX’s auto-generated documentation is athttps://docs.jax.dev/.

The documentation building is controlled for the entire project by thereadthedocs JAX settings. The current settingstrigger a documentation build as soon as code is pushed to the GitHubmain branch.For each code version, the building process is driven by the.readthedocs.yml and thedocs/conf.py configuration files.

For each automated documentation build you can see thedocumentation build logs.

If you want to test the documentation generation on Readthedocs,you can add the “documentation” GitHub label to your PR. You will thenbe able to view the docs from the link for the “docs/readthedocs.org:jax”GitHub check.

For a local test, I was able to do it in a fresh directory by replaying the commandsI saw in the Readthedocs logs:

mkvirtualenv jax-docs  # A new virtualenvmkdir jax-docs  # A new directorycd jax-docsgit clone --no-single-branch --depth 50 https://github.com/jax-ml/jaxcd jaxgit checkout --force origin/test-docsgit clean -d -f -fworkon jax-docspython -m pip install --upgrade --no-cache-dir pippython -m pip install --upgrade --no-cache-dir -I Pygments==2.3.1 setuptools==41.0.1 docutils==0.14 mock==1.0.1 pillow==5.4.1 alabaster>=0.7,<0.8,!=0.7.5 commonmark==0.8.1 recommonmark==0.5.0 'sphinx<2' 'sphinx-rtd-theme<0.5' 'readthedocs-sphinx-ext<1.1'python -m pip install --exists-action=w --no-cache-dir -r docs/requirements.txtcd docspython `which sphinx-build` -T -E -b html -d _build/doctrees-readthedocs -D language=en . _build/html