## File: README.md # Transformable numerical computing at scale [](https://github.com/jax-ml/jax/actions/workflows/ci-build.yaml) [](https://pypi.org/project/jax/) [**Transformations**](#transformations) | [**Scaling**](#scaling) | [**Install guide**](#installation) | [**Change logs**](https://docs.jax.dev/en/latest/changelog.html) | [**Reference docs**](https://docs.jax.dev/en/latest/) ## What is JAX? JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning. JAX can automatically differentiate native Python and NumPy functions. It can differentiate through loops, branches, recursion, and closures, and it can take derivatives of derivatives of derivatives. It supports reverse-mode differentiation (a.k.a. backpropagation) via [`jax.grad`](#automatic-differentiation-with-grad) as well as forward-mode differentiation, and the two can be composed arbitrarily to any order. JAX uses [XLA](https://www.openxla.org/xla) to compile and scale your NumPy programs on TPUs, GPUs, and other hardware accelerators. You can compile your own pure functions with [`jax.jit`](#compilation-with-jit). Compilation and automatic differentiation can be composed arbitrarily. Dig a little deeper, and you'll see that JAX is really an extensible system for [composable function transformations](#transformations) at [scale](#scaling). This is a research project, not an official Google product. Expect [sharp edges](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html). Please help by trying it out, [reporting bugs](https://github.com/jax-ml/jax/issues), and letting us know what you think! ```python import jax import jax.numpy as jnp def predict(params, inputs): for W, b in params: outputs = jnp.dot(inputs, W) + b inputs = jnp.tanh(outputs) # inputs to the next layer return outputs # no activation on last layer def loss(params, inputs, targets): preds = predict(params, inputs) return jnp.sum((preds - targets)**2) grad_loss = jax.jit(jax.grad(loss)) # compiled gradient evaluation function perex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0))) # fast per-example grads ``` ### Contents * [Transformations](#transformations) * [Scaling](#scaling) * [Current gotchas](#gotchas-and-sharp-bits) * [Installation](#installation) * [Citing JAX](#citing-jax) * [Reference documentation](#reference-documentation) ## Transformations At its core, JAX is an extensible system for transforming numerical functions. Here are three: `jax.grad`, `jax.jit`, and `jax.vmap`. ### Automatic differentiation with `grad` Use [`jax.grad`](https://docs.jax.dev/en/latest/jax.html#jax.grad) to efficiently compute reverse-mode gradients: ```python import jax import jax.numpy as jnp def tanh(x): y = jnp.exp(-2.0 * x) return (1.0 - y) / (1.0 + y) grad_tanh = jax.grad(tanh) print(grad_tanh(1.0)) # prints 0.4199743 ``` You can differentiate to any order with `grad`: ```python print(jax.grad(jax.grad(jax.grad(tanh)))(1.0)) # prints 0.62162673 ``` You're free to use differentiation with Python control flow: ```python def abs_val(x): if x > 0: return x else: return -x abs_val_grad = jax.grad(abs_val) print(abs_val_grad(1.0)) # prints 1.0 print(abs_val_grad(-1.0)) # prints -1.0 (abs_val is re-evaluated) ``` See the [JAX Autodiff Cookbook](https://docs.jax.dev/en/latest/notebooks/autodiff_cookbook.html) and the [reference docs on automatic differentiation](https://docs.jax.dev/en/latest/jax.html#automatic-differentiation) for more. ### Compilation with `jit` Use XLA to compile your functions end-to-end with [`jit`](https://docs.jax.dev/en/latest/jax.html#just-in-time-compilation-jit), used either as an `@jit` decorator or as a higher-order function. ```python import jax import jax.numpy as jnp def slow_f(x): # Element-wise ops see a large benefit from fusion return x * x + x * 2.0 x = jnp.ones((5000, 5000)) fast_f = jax.jit(slow_f) %timeit -n10 -r3 fast_f(x) %timeit -n10 -r3 slow_f(x) ``` Using `jax.jit` constrains the kind of Python control flow the function can use; see the tutorial on [Control Flow and Logical Operators with JIT](https://docs.jax.dev/en/latest/control-flow.html) for more. ### Auto-vectorization with `vmap` [`vmap`](https://docs.jax.dev/en/latest/jax.html#vectorization-vmap) maps a function along array axes. But instead of just looping over function applications, it pushes the loop down onto the function’s primitive operations, e.g. turning matrix-vector multiplies into matrix-matrix multiplies for better performance. Using `vmap` can save you from having to carry around batch dimensions in your code: ```python import jax import jax.numpy as jnp def l1_distance(x, y): assert x.ndim == y.ndim == 1 # only works on 1D inputs return jnp.sum(jnp.abs(x - y)) def pairwise_distances(dist1D, xs): return jax.vmap(jax.vmap(dist1D, (0, None)), (None, 0))(xs, xs) xs = jax.random.normal(jax.random.key(0), (100, 3)) dists = pairwise_distances(l1_distance, xs) dists.shape # (100, 100) ``` By composing `jax.vmap` with `jax.grad` and `jax.jit`, we can get efficient Jacobian matrices, or per-example gradients: ```python per_example_grads = jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0, 0))) ``` ## Scaling To scale your computations across thousands of devices, you can use any composition of these: * [**Compiler-based automatic parallelization**](https://docs.jax.dev/en/latest/parallel.html) where you program as if using a single global machine, and the compiler chooses how to shard data and partition computation (with some user-provided constraints); * [**Explicit sharding and automatic partitioning**](https://docs.jax.dev/en/latest/parallel.html) where you still have a global view but data shardings are explicit in JAX types, inspectable using `jax.typeof`; * [**Manual per-device programming**](https://docs.jax.dev/en/latest/notebooks/shard_map.html) where you have a per-device view of data and computation, and can communicate with explicit collectives. | Mode | View? | Explicit sharding? | Explicit Collectives? | |---|---|---|---| | Auto | Global | ❌ | ❌ | | Explicit | Global | ✅ | ❌ | | Manual | Per-device | ✅ | ✅ | ```python from jax.sharding import set_mesh, AxisType, PartitionSpec as P mesh = jax.make_mesh((8,), ('data',), axis_types=(AxisType.Explicit,)) set_mesh(mesh) # parameters are sharded for FSDP: for W, b in params: print(f'{jax.typeof(W)}') # f32[512@data,512] print(f'{jax.typeof(b)}') # f32[512] # shard data for batch parallelism: inputs, targets = jax.device_put((inputs, targets), P('data')) # evaluate gradients, automatically parallelized! gradfun = jax.jit(jax.grad(loss)) param_grads = gradfun(params, (inputs, targets)) ``` See the [tutorial](https://docs.jax.dev/en/latest/parallel.html) and [advanced guides](https://docs.jax.dev/en/latest/advanced_guide.html) for more. ## Gotchas and sharp bits See the [Gotchas Notebook](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html). ## Installation ### Supported platforms | | Linux x86_64 | Linux aarch64 | Mac aarch64 | Windows x86_64 | Windows WSL2 x86_64 | |------------|--------------|---------------|--------------|----------------|---------------------| | CPU | yes | yes | yes | yes | yes | | NVIDIA GPU | yes | yes | n/a | no | experimental | | Google TPU | yes | n/a | n/a | n/a | n/a | | AMD GPU | yes | no | n/a | no | experimental | | Apple GPU | n/a | no | experimental | n/a | n/a | | Intel GPU | experimental | n/a | n/a | no | no | ### Instructions | Platform | Instructions | |-----------------|-----------------------------------------------------------------------------------------------------------------| | CPU | `pip install -U jax` | | NVIDIA GPU | `pip install -U "jax[cuda13]"` | | Google TPU | `pip install -U "jax[tpu]"` | | AMD GPU (Linux) | `pip install -U "jax[rocm7-local]"`| | Intel GPU | Follow [Intel's instructions](https://github.com/intel/intel-extension-for-openxla/blob/main/docs/acc_jax.md). | See [the documentation](https://docs.jax.dev/en/latest/installation.html) for information on alternative installation strategies. These include compiling from source, installing with Docker, using other versions of CUDA, a community-supported conda build, and answers to some frequently-asked questions. ## Citing JAX To cite this repository: ``` @software{jax2018github, author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Yash Katariya and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang}, title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs}, url = {http://github.com/jax-ml/jax}, version = {0.3.13}, year = {2018}, } ``` In the above bibtex entry, names are in alphabetical order, the version number is intended to be that from [jax/version.py](../main/jax/version.py), and the year corresponds to the project's open-source release. A nascent version of JAX, supporting only automatic differentiation and compilation to XLA, was described in a [paper that appeared at SysML 2018](https://mlsys.org/Conferences/2019/doc/2018/146.pdf). We're currently working on covering JAX's ideas and capabilities in a more comprehensive and up-to-date paper. ## Reference documentation For details about the JAX API, see the [reference documentation](https://docs.jax.dev/). For getting started as a JAX developer, see the [developer documentation](https://docs.jax.dev/en/latest/developer.html). --- ## File: ci/envs/README.md # JAXCI Environment Variables This docpage describes the various `JAXCI` environment variables that are used in the CI scripts and their behaviors. These variables are used to control the behavior of the CI scripts such as the Python version used, path to JAX/XLA repo, if to clone XLA repo, etc. Name | Default Value | Behavior | Usage ------------------------------------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ----- `JAXCI_JAX_GIT_DIR` | Present working directory: `$(pwd)` | Path to the JAX's Git directory. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_JAX_GIT_DIR&type=code) `JAXCI_HERMETIC_PYTHON_VERSION` | System default | Controls the version of hermetic Python to use. This affects the Bazel commands only such as when building artifacts or when running the Bazel test scripts. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_HERMETIC_PYTHON_VERSION&type=code) `JAXCI_CUDA_VERSION` | 12 | Controls the CUDA version to use when building the JAX artifacts or running the tests. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_CUDA_VERSION&type=code) `JAXCI_ROCM_VERSION` | 7 | Controls the ROCm version to use when building the JAX artifacts or running the tests. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_ROCM_VERSION&type=code) `JAXCI_XLA_GIT_DIR` | Unset | When using a local copy of XLA, this points to the root of the XLA git repository. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_XLA_GIT_DIR&type=code) `JAXCI_CLONE_MAIN_XLA` | 0 | If set to 1, the XLA repository is cloned at HEAD and its path is set in `JAXCI_XLA_GIT_DIR` | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_CLONE_MAIN_XLA&type=code) `JAXCI_XLA_COMMIT` | Unset | Allows overriding the XLA commit that is used when using a local copy of XLA. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_XLA_COMMIT&type=code) `JAXCI_OUTPUT_DIR` | `$(pwd)/dist` | Controls the location where the artifacts are written to. The directory will be automatically created if it does not exist. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_OUTPUT_DIR&type=code) `JAXCI_BUILD_ARTIFACT_WITH_RBE` | 0 | When set to 1, Bazel will use RBE to build the artifacts. Requires gcloud authentication and only certain platforms support RBE so this typically only set in CI builds | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_BUILD_ARTIFACT_WITH_RBE&type=code) `JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE` | 0 | When set to 1, Bazel will also try to push new cache entries to the cache bucket. Since writes to the bucket require authentication, this flag is enabled only for CI builds. Note that the builds using RBE use the RBE cache and not Bazel's remote cache, therefore this variable is a no-op if `JAXCI_BUILD_ARTIFACT_WITH_RBE` is set to 1. When `JAXCI_BUILD_ARTIFACT_WITH_RBE` and `JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE` are both not set, Bazel will still read from the public cache bucket to try to speed up the build. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_WRITE_TO_BAZEL_REMOTE_CACHE&type=code) `JAXCI_ARTIFACT_TYPE` | "default" | Controls the type of artifacts to build. Valid values are "default", "release", "nightly". This affects the wheel tag and metadata, see [ci/build_artifacts.sh](https://github.com/jax-ml/jax/blob/main/ci/build_artifacts.sh) to understand how. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_ARTIFACT_TYPE&type=code) `JAXCI_WHEEL_RC_VERSION` | Unset | During the release process, we build a Release Candidate (RC) wheel in addition to the release wheel. This environment variable sets the version of the RC wheel to build. Values are set internally. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_WHEEL_RC_VERSION&type=code) `JAXCI_PYTHON` | `python${JAXCI_HERMETIC_PYTHON_VERSION}` | Points to the system Python binary to use. It used by scripts that make use of the system Python such as the Pytest scripts. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_PYTHON&type=code) `JAXCI_ENABLE_X64` | 0 | By default, JAX enforces single-precision numbers to mitigate the Numpy API’s tendency to aggressively promote operands to `double`. When set to 1, the tests will use double-precision numbers. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_ENABLE_X64&type=code) `JAXCI_TPU_CORES` | Unset | Sets the number of TPU cores for the TPU machine type. Values are set in the workflow files. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_TPU_CORES&type=code) `JAXCI_RUN_FULL_TPU_TEST_SUITE` | 0 | When set to 1, the full TPU test suite is run. Otherwise, a subset of tests is run. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_RUN_FULL_TPU_TEST_SUITE&type=code) `JAXCI_JAX_PYPI_EXTRAS` | Unset | Used to control the installation of JAX extras from PyPI. See JAX's [setup.py](https://github.com/jax-ml/jax/blob/c9934912885bb7c4b72c5a9271598235a6789a81/setup.py#L71) for the list of valid values. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_JAX_PYPI_EXTRAS&type=code) `JAXCI_BUILD_JAXLIB` | true | Used to control the value of [build_jaxlib](https://github.com/jax-ml/jax/blob/338b4ebc8a5478e3d22efc9530be71d69c3bb993/jax/BUILD#L55-L63) flag. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_BUILD_JAXLIB&type=code) `JAXCI_BUILD_JAX` | true | Used to control the value of [build_jax](https://github.com/jax-ml/jax/blob/338b4ebc8a5478e3d22efc9530be71d69c3bb993/jax/BUILD#L92-L100) flag. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_BUILD_JAX&type=code) `JAXCI_BAZEL_OUTPUT_BASE` | Unset | Used to control the output base for Bazel builds. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_BAZEL_OUTPUT_BASE&type=code) `JAXCI_BAZEL_CPU_RBE_MODE` | test | Used to control whether to run or build the CPU test targets. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_BAZEL_CPU_RBE_MODE&type=code) ## Docker Specific Environment Variables > [!NOTE] > The following environment variables only affect the build if the > [run_docker_container.sh](https://github.com/jax-ml/jax/blob/main/ci/utilities/run_docker_container.sh) > script was invoked to start a Docker container and the build is running inside > that container. Typically, this would be the internal CI builds and local > builds. Note that while GitHub actions use the same Docker images, they do not > invoke "run_docker_container.sh" as they leverage built-in containerization > features to run jobs within a container. Name | Default Value | Behavior | Usage ----------------------- | ------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------- | ----- `JAXCI_DOCKER_WORK_DIR` | "/jax" | The path on the container where the JAX Git repository is mounted to. | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_DOCKER_WORK_DIR&type=code) `JAXCI_DOCKER_ARGS` | Empty String | Space separated string of additional arguments that will be passed when starting the Docker container | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_DOCKER_ARGS&type=code) `JAXCI_DOCKER_IMAGE` | Depends on the system (see [ci/envs/docker.env](https://github.com/jax-ml/jax/blob/main/ci/envs/docker.env)) | Docker image to pull | [Usage](https://github.com/search?q=repo%3Ajax-ml%2Fjax%20JAXCI_DOCKER_IMAGE&type=code) --- ## File: ci/utilities/README.md # JAX CI Utility Scripts This docpage gives a brief overview of the different utility scripts and what they are used for. - **setup_build_environment.sh**: Sets up the build environment such as cloning the latest XLA, adjusting file paths (for Windows), etc. - **convert_msys_paths_to_win_paths.py**: Converts MSYS Linux-like paths stored in env variables to Windows paths. - **install_wheels_locally.sh**: Used by Pytest scripts to install JAX wheels and any additional extras on the system. - **run_auditwheel.sh**: Verifies that the Linux artifacts are "manylinux" compliant. - **run_docker_container.sh**: Runs a Docker container called "jax". Images are read from the `JAXCI_DOCKER_IMAGE` environment variable in [ci/envs/docker.env](https://github.com/jax-ml/jax/blob/main/ci/envs/docker.env). --- ## File: ci/README.md # JAX Continuous Integration This folder contains the configuration files and scripts used to build and test JAX. It is typically used by continuous integration (CI) jobs to automate builds and run comprehensive tests across various platforms and configurations. This page provides an overview of the JAX CI system, its components, and the different workflows it supports. ******************************************************************************** ## JAX's CI System JAX's CI system is composed of several interacting components and orchestrates builds and tests using a hybrid approach, leveraging both an internal CI system and GitHub Actions as well as an internal build orchestrator for managing nightly and release flows. It encompasses several distinct workflows, including comprehensive presubmit checks triggered on pull requests and branch pushes, bi-hourly continuous builds, extensive nightly builds with broad platform coverage, and a controlled release process that culminates in PyPI publication. These flows build four packages: `jax`, `jaxlib`, `jax-cuda-plugin`, `jax-cuda-pjrt` and support a range of environments, including: * **Linux x86:** CPU, TPU, CUDA * **Linux aarch64:** CPU, CUDA * **Windows x86:** CPU * **Mac Arm64:** CPU ### Architecture Overview 1. **Internal CI System:** An internal CI system is used for specific build and test tasks, such as nightly builds, release candidate (RC) builds, and Mac-specific testing. 2. **GitHub Actions:** Used for presubmit checks, continuous integration builds and tests, and nightly/release artifact testing. 3. **Build Orchestrator:** An internal tool used to manage complex workflows such as nightly / release flows, promoting RC builds to release, etc. 4. **Artifact Storage:** * Google Cloud Storage (GCS) Buckets: Used for temporary storage of artifacts between jobs in GitHub Actions workflows and for storing packages built during nightly and release flows before testing. * Artifact Registry: Used to store nightly packages, RC packages and final releases. * PyPI: Where final releases are published. ### CI Workflows and Where They Run JAX's CI system consists of the following workflows: 1. **Presubmits:** Presubmits are run in GitHub actions and are triggered on pull requests that target the `main` branch and on pushes to the `main` and `release` branch. JAX's presubmit run time SLO is about 10 minutes so these are typically run using Bazel with remote build execution ([RBE](https://bazel.build/remote/rbe)). RBE allows us to execute build and test actions on a distributed system, separate from the local machine, instead of solely on the local machine. This enables faster build and test times by utilizing parallel computing resources and caching across a cluster of machines. However, we also use Pytest in workflows where we are not able to use RBE such as the TPU presubmit. In such presubmits, we usually run a subset of tests to be able to satisfy the presubmit run time SLO. To see the list of the presubmit workflows, [click here](https://github.com/search?q=repo%3Ajax-ml%2Fjax+path%3A.github%2Fworkflows%2F+%28path%3A**%2F*.yml+OR+path%3A**%2F*.yaml%29+%22pull_request%22&type=code). 2. **Continuous:** These jobs are run in GitHub actions and are scheduled to run once every 2 hours on the `main` branch. It builds JAX packages and runs a wide range of tests targeting different environments such as CPU, CUDA (L4, H100, B200, etc), and TPU (v4-8, v5e-8, etc.). For more information, see [wheel_tests_continuous.yml](https://github.com/jax-ml/jax/blob/main/.github/workflows/wheel_tests_continuous.yml) ([An example run](https://github.com/jax-ml/jax/actions/workflows/wheel_tests_continuous.yml).) 3. **Nightly Builds and Tests:** These jobs use an hybrid approach of both the internal CI system and GitHub actions. The jobs are triggered once every night by the internal build orchestrator tool. It first triggers the jobs in the internal CI system to build the JAX packages for different configurations (Python versions, CUDA versions, etc) and uploads them to a staging bucket in GCS as well as to the nightly artifact registry. Next, testing jobs are triggered that download the artifacts from the staging bucket and run tests. Mac testing jobs are run in the internal CI system. For non-Mac testing, a trigger job is run that invokes the [wheel_tests_nightly_release.yml](https://github.com/jax-ml/jax/blob/main/.github/workflows/wheel_tests_nightly_release.yml) workflow in GitHub Actions. JAX's nightly artifacts can be found here: [jax](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jax), [jaxlib](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jaxlib), [jax-cuda-plugin](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jax-cuda12-plugin), [jax-cuda-pjrt](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-nightly-artifacts-registry/simple/jax-cuda12-pjrt). 4. **Release Builds and Tests:** Release flow is similar to the nightly flow except for few differences. First, release process has to be triggered manually in the internal build orchestrator and should be done only after a release branch (E.g `release/0.5.3`) has been created. The build jobs build two sets of artifacts for each package: 1. RC wheels 2. Final version wheels. These two sets are pretty much the same package except for their metadata and wheel tags. The RC wheels are then uploaded to the staging bucket and release artifact registry. After the uploads are done, the test jobs are triggered. As with the nightly flow, Mac test jobs are run in the internal CI system while non-Mac test jobs are run in GitHub actions. To see the GitHub actions run for a particular release, filter the workflow runs by its branch name. 5. **Promote RC to Final and Publish to PyPI:** If the RC wheels pass all testing, then we are ready to promote it as the final version and publish it to PyPI. This entire flow is internal and is run in our internal CI system. Final version of the packages are published to PyPI and JAX's release artifact registry. JAX's release artifacts (RC and final versions) can be found here: [jax](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-release-artifacts-registry/simple/jax), [jaxlib](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-release-artifacts-registry/simple/jaxlib), [jax-cuda-plugin](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-release-artifacts-registry/simple/jax-cuda12-plugin), [jax-cuda-pjrt](https://us-python.pkg.dev/ml-oss-artifacts-published/jax-public-release-artifacts-registry/simple/jax-cuda12-pjrt). ### JAX's Official CI and Build/Test Scripts JAX's CI jobs (both internal and those on GitHub actions) run the scripts in this folder. An overview of the different folders and their purpose is given below: - **ci/**: Contains all build scripts, environment files, and utility scripts. - **ci/utilities/**: Contains helper scripts used throughout the build/test process. See [README.md](https://github.com/jax-ml/jax/blob/main/ci/utilities/README.md) for a brief overview of these utility scripts and their behavior. - **ci/envs/**: Holds environment files that set `JAXCI` environment variables that control build and test configurations. see [README.md](https://github.com/jax-ml/jax/blob/main/ci/envs/README.md) to see the complete list of these variables and their behavior. Every build script in this folder first source the `JAXCI` envs in [default.env](https://github.com/jax-ml/jax/blob/main/ci/envs/default.env) and then run the [setup_build_environment.sh](https://github.com/jax-ml/jax/blob/main/ci/utilities/setup_build_environment.sh) script to set up the build environment. A brief overview of each build script in this folder is given below: > [!NOTE] > Both internal and GitHub action jobs run under the > [ml-build](https://github.com/tensorflow/tensorflow/tree/master/ci/official/containers) > Docker image which contains build tools such as Python, Bazelisk, LLVM/Clang, > manylinux compliant libraries (in Linux images), etc. - **build_artifacts.sh:** These build the various JAX artifacts. We build three different type of artifacts based on the type of job: Nightly, RC/Release, or at HEAD. - **run_bazel_test_cpu_rbe.sh/run_bazel_test_cuda_rbe.sh**: These run Bazel tests with RBE on every GitHub PR. We test compatibility with both CPU and CUDA. On platforms where RBE is not natively supported (e.g Linux Arm64), we cross-compile the test targets for Linux Aarch64 on Linux x86. As the tests still need to be run on the host machines and because running the tests on a single machine can take a long time, we skip running them on these platforms. Note for `run_bazel_test_cpu_rbe.sh`: - If `$JAXCI_BUILD_JAXLIB=false` and `$JAXCI_BUILD_JAX=false`, these jobs depend on local JAX wheels and therefore require that the following wheels to be present in the `../dist` folder: `jax`, and `jaxlib` wheels. In CI builds, we first build these wheels from source and then run the `bazel test` command. - If `$JAXCI_BUILD_JAXLIB=false` and `$JAXCI_BUILD_JAX=true`, CPU jobs depend on local jaxlib wheels and therefore require that `jaxlib` wheel to be present in the `../dist` folder. GPU obs depend on local jaxlib and CUDA wheels, and therefore require that the following wheels to be present in the `../dist` folder: `jaxlib`, `jax-cuda-plugin`, and `jax-cuda-pjrt` wheels. In CI builds, we first build these wheels from source and then run the `bazel test` command. - If `$JAXCI_BUILD_JAXLIB=wheel` and `$JAXCI_BUILD_JAX=wheel`, the Bazel tests use [py_import](https://github.com/openxla/xla/blob/8190847008eddd4c7f3e57449e16d28631770823/third_party/py/py_import.bzl#L47). - If `$JAXCI_BUILD_JAXLIB=true` and `$JAXCI_BUILD_JAX=true`, Bazel will use individual targets in the test dependencies. - **run_bazel_test_cuda_non_rbe.sh**: These run the following Bazel CUDA tests: Single accelerator tests with one GPU apiece and Multi-accelerator tests with all GPUs. - If `$JAXCI_BUILD_JAXLIB=false` and `$JAXCI_BUILD_JAX=false`, these jobs depend on local JAX wheels and therefore require that the following wheels to be present in the `../dist` folder: `jax`, `jaxlib`, `jax-cuda-plugin`, and `jax-cuda-pjrt` wheels. In CI builds, we first build these wheels from source and then run the `bazel test` command. - If `$JAXCI_BUILD_JAXLIB=wheel` and `$JAXCI_BUILD_JAX=wheel`, the Bazel tests use [py_import](https://github.com/openxla/xla/blob/8190847008eddd4c7f3e57449e16d28631770823/third_party/py/py_import.bzl#L47). - **run_pytest_*.sh**: These run tests with Pytests and use the JAX wheel packages installed on the system. In CI builds, we build the wheels first from source and then run the `pytest` commands. We test compatibility with CPU, CUDA, and TPU. These are primarily run as part of the continuous and nightly/release test jobs except for TPU which is also run as a presubmit testing a subset of the tests. ## Different Test Configurations JAX's CI Test jobs run under different test configurations. These configurations are described briefly in the sections below. ### XLA Versions JAX's CI builds rely on XLA, but use different versions depending on the type of build. To ensure stability and reproducibility, nightly and release builds use a pinned XLA version defined in [MODULE.bazel](https://github.com/jax-ml/jax/blob/main/MODULE.bazel). However, to keep JAX compatible with the latest XLA developments, presubmit and postsubmit builds utilize the most recent XLA version. This is done by overriding the default XLA dependency with a local copy of the XLA repository. We do this by passing `--override_module=xla=/path/to/local/xla` which instructs Bazel to depend on the XLA in the local system instead of the pinned version in `MODULE.bazel`. The CI system uses the `JAXCI` environment variables to manage this process. When running jobs that need to use XLA at head, we set `JAXCI_CLONE_MAIN_XLA=1`. This clones the XLA repository at head and sets `JAXCI_XLA_GIT_DIR` to its path. [JAX build CLI](https://github.com/jax-ml/jax/blob/main/build/build.py) automatically adds the necessary Bazel flag (`--override_repository`) to point to this local XLA version during the build process if `JAXCI_XLA_GIT_DIR` is set. In jobs where the build CLI is not used such as the RBE presubmits, we explicitly include `--override_repository=xla="${JAXCI_XLA_GIT_DIR}"` as part of the test command. ### Enabling/Disabling 64-bit Data Types By default, JAX enforces single-precision numbers to mitigate the Numpy API’s tendency to aggressively promote operands to `double`. In order to use double-precision numbers, we need to set the `JAX_ENABLE_X64` environment variable. In CI, we test both configurations in presubmits and postsubmits by using the `JAXCI_ENABLE_X64` environment variable. ## [Googlers Only] Connecting to CI Runners for Debugging If you are a Googler, you can connect to one of the self-hosted runners we have on GitHub to debug your workflow. For more information, see go/ml-github-actions:connect. ## Running These Scripts Locally on Your Machine > [!IMPORTANT] > If you are a Linux / Windows user, you need to have Docker installed as a > prerequisite. Additionally, if running on Windows, please run these commands > in a bash environment as all the scripts are written in Shell. Follow the steps below to run a CI script locally on your machine. 1. [Optional] Set `JAXCI` variables in your shell environment. See [ci/envs/README.md](https://github.com/jax-ml/jax/blob/main/ci/envs/README.md) for the list of `JAXCI` variables and their behavior. 2. [Linux/Windows] Start the Docker container by running: ```bash ./ci/utilities/run_docker_container.sh ``` This will start a Docker container named "jax". Note that if you set any `JAXCI` variables in step 1, they will also be be set in the container. Run the script under the Docker container. ```bash # docker exec jax docker exec jax ./ci/build_artifacts.sh jaxlib ``` 3. [Mac] Execute the build script directly. ```bash # ./ ./ci/build_artifacts.sh jaxlib ``` --- ## File: cloud_tpu_colabs/README.md # JAX on Cloud TPU examples The same JAX code that runs on CPU and GPU can also be run on TPU. Cloud TPUs have the advantage of quickly giving you access to multiple TPU accelerators, including in [Colab](https://research.google.com/colaboratory/). All of the example notebooks here use [`jax.pmap`](https://docs.jax.dev/en/latest/jax.html#jax.pmap) to run JAX computation across multiple TPU cores from Colab. You can also run the same code directly on a [Cloud TPU VM](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm). ## Example Cloud TPU notebooks The following notebooks showcase how to use and what you can do with Cloud TPUs on Colab: ### [Pmap Cookbook](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/Pmap_Cookbook.ipynb) A guide to getting started with `pmap`, a transform for easily distributing SPMD computations across devices. ### [Lorentz ODE Solver](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/Lorentz_ODE_Solver.ipynb) Contributed by Alex Alemi (alexalemi@) Solve and plot parallel ODE solutions with `pmap`. ### [Wave Equation](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/Wave_Equation.ipynb) Contributed by Stephan Hoyer (shoyer@) Solve the wave equation with `pmap`, and make cool movies! The spatial domain is partitioned across the 8 cores of a Cloud TPU. ### [JAX Demo](https://colab.research.google.com/github/jax-ml/jax/blob/main/cloud_tpu_colabs/JAX_demo.ipynb) An overview of JAX presented at the [Program Transformations for ML workshop at NeurIPS 2019](https://program-transformations.github.io/) and the [Compilers for ML workshop at CGO 2020](https://www.c4ml.org/). Covers basic numpy usage, `grad`, `jit`, `vmap`, and `pmap`. ## Performance notes The [guidance on running TensorFlow on TPUs](https://cloud.google.com/tpu/docs/performance-guide) applies to JAX as well, with the exception of TensorFlow-specific details. Here we highlight a few important details that are particularly relevant to using TPUs in JAX. ### Padding One of the most common culprits for surprisingly slow code on TPUs is inadvertent padding: - Arrays in the Cloud TPU are tiled. This entails padding one of the dimensions to a multiple of 8, and a different dimension to a multiple of 128. - The matrix multiplication unit performs best with pairs of large matrices that minimize the need for padding. ### bfloat16 dtype By default\*, matrix multiplication in JAX on TPUs [uses bfloat16](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus) with float32 accumulation. This can be controlled with the `precision` keyword argument on relevant `jax.numpy` functions (`matmul`, `dot`, `einsum`, etc). In particular: - `precision=jax.lax.Precision.DEFAULT`: uses mixed bfloat16 precision (fastest) - `precision=jax.lax.Precision.HIGH`: uses multiple MXU passes to achieve higher precision - `precision=jax.lax.Precision.HIGHEST`: uses even more MXU passes to achieve full float32 precision JAX also adds the `bfloat16` dtype, which you can use to explicitly cast arrays to bfloat16, e.g., `jax.numpy.array(x, dtype=jax.numpy.bfloat16)`. \* We might change the default precision in the future, since it is arguably surprising. Please comment/vote on [this issue](https://github.com/jax-ml/jax/issues/2161) if it affects you! ## Running JAX on a Cloud TPU VM Refer to the [Cloud TPU VM documentation](https://cloud.google.com/tpu/docs/jax-quickstart-tpu-vm). ## Reporting issues and getting help If you run into Cloud TPU-specific issues (e.g. trouble creating a Cloud TPU VM), please email , or if you are a [TRC](https://sites.research.google/trc/) member. You can also [file a JAX issue](https://github.com/jax-ml/jax/issues) or [ask a discussion question](https://github.com/jax-ml/jax/discussions) for any issues with these notebooks or using JAX in general. If you have any other questions or comments regarding JAX on Cloud TPUs, please email . We’d like to hear from you! --- ## File: docs/_tutorials/advanced-compilation.md # Advanced compilation ```{note} This is a placeholder for a section in the new {ref}`jax-tutorials-draft`. For the time being, you may find some related content in the old documentation: - {doc}`../aot` - {doc}`../pallas/index` ``` --- ## File: docs/_tutorials/advanced-debugging.md --- jupytext: formats: md:myst text_representation: extension: .md format_name: myst format_version: 0.13 jupytext_version: 1.16.4 kernelspec: display_name: Python 3 language: python name: python3 --- (advanced-debugging)= # Advanced debugging ```{note} This is a placeholder for a section in the new {ref}`jax-tutorials-draft`. For the time being, you may find some related content in the old documentation: - {doc}`../debugging` ``` --- ## File: docs/_tutorials/index.rst :orphan: .. _jax-tutorials-draft: JAX tutorials draft =================== .. note:: The tutorials below are a work in progress; for the time being, please refer to the older tutorial content, including :ref:`beginner-guide`, :ref:`jax-101`, and the now-deleted *JAX 101* tutorials. JAX 101 ------- Mostly finalized at :ref:`jax-101`! .. toctree:: :maxdepth: 1 ../key-concepts ../jit-compilation ../automatic-vectorization ../automatic-differentiation ../debugging ../random-numbers ../pytrees ../parallel ../stateful-computations simple-neural-network JAX 201 ------- .. toctree:: :maxdepth: 1 parallelism advanced-debugging profiling-and-performance JAX 301 ------- .. toctree:: :maxdepth: 1 advanced-compilation --- ## File: docs/_tutorials/parallelism.md # Parallel computation ```{note} This is a placeholder for a section in the new {ref}`jax-tutorials-draft`. For the time being, you may find some related content in the old documentation: - {doc}`../multi_process` - {doc}`../parallel` ``` --- ## File: docs/_tutorials/profiling-and-performance.md # Profiling and performance ```{note} This is a placeholder for a section in the new {ref}`jax-tutorials-draft`. For the time being, you may find some related content in the old documentation: - {doc}`../profiling` - {doc}`../device_memory_profiling` - {doc}`../transfer_guard` ```