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 as well as forward-mode differentiation, and the two can be composed arbitrarily to any order. JAX uses 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 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 at scale. See the JAX Autodiff Cookbook and the reference docs on automatic differentiation for more. 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 for more. 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. By composing jax.vmap with jax.grad and jax.jit , we can get efficient Jacobian matrices, or per-example gradients: To scale your computations across thousands of devices, you can use any composition of these: Compiler-based automatic parallelization 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 where you still have a global view but data shardings are explicit in JAX types, inspectable using jax.typeof ; Manual per-device programming where you have a per-device view of data and computation, and can communicate with explicit collectives. 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. See the documentation 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. In the above bibtex entry, names are in alphabetical order, the version number is intended to be that from 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. We're currently working on covering JAX's ideas and capabilities in a more comprehensive and up-to-date paper.