JAX

JAX is an open-source Python library developed under the jax-ml GitHub organization, originally created by researchers including James Bradbury, Roy Frostig, and Adam Paszke. It provides composable function transformations — differentiation, vectorization, JIT compilation, and parallelization — on top of a NumPy-compatible API. The project describes itself as a research project, not an official Google product, and notes users should expect sharp edges.

What It Is

JAX sits at the intersection of numerical computing and machine learning infrastructure. It gives Python developers a familiar NumPy-style API while enabling the same code to run efficiently on CPUs, NVIDIA GPUs, Google TPUs, AMD GPUs, and Intel GPUs. Under the hood, JAX uses XLA (the Accelerated Linear Algebra compiler from OpenXLA) to compile and optimize programs for hardware accelerators. The core design philosophy is composability: a small set of powerful transformations can be combined arbitrarily to build complex workflows.

Core Transformations

JAX's value proposition centers on three composable primitives:

• jax.grad: Automatic differentiation supporting both reverse-mode (backpropagation) and forward-mode differentiation. It can differentiate through loops, branches, recursion, and closures, and supports higher-order derivatives (derivatives of derivatives).
• jax.jit: Just-in-time compilation via XLA, which fuses element-wise operations and compiles pure Python/NumPy functions end-to-end for significant speedups on accelerators.
• jax.vmap: Automatic vectorization that maps a function along array axes by pushing the batch loop down into primitive operations (e.g., turning matrix-vector multiplies into matrix-matrix multiplies), eliminating the need to manually manage batch dimensions.

These three can be freely composed — for example, jax.jit(jax.vmap(jax.grad(loss))) produces a compiled, batched gradient function in a single expression.

Scaling to Multi-Device Workloads

JAX provides three modes for distributing computation across many devices:

• Auto mode: Program as if using a single global machine; the compiler chooses how to shard data and partition computation.
• Explicit mode: Global view with explicit data shardings visible in JAX types, inspectable via jax.typeof.
• Manual mode: Per-device view with explicit collective communication primitives.

This flexibility supports use cases from single-GPU research experiments to large-scale distributed training across thousands of accelerators.

Ecosystem and Integrations

JAX itself is narrowly scoped to array operations and program transformations. A broad ecosystem of libraries builds on top of it, including:

• Neural networks: Flax, Equinox, Keras
• Optimizers & solvers: Optax, Optimistix, Diffrax
• Data loading: Grain, TensorFlow Datasets, Hugging Face Datasets
• Probabilistic programming: Blackjax, Numpyro, PyMC
• LLM training: MaxText, AXLearn, Levanter, Marin
• Physics & simulation: JAX MD, Brax

The community-maintained Awesome JAX list tracks the broader ecosystem.

Update: JAX v0.10.1

The latest release is JAX v0.10.1, published on May 20, 2026. The repository remains actively maintained with frequent pushes and over 35,000 GitHub stars and 3,500 forks as of mid-2026. The project is licensed under Apache License 2.0. Installation is straightforward via pip, with separate packages for CPU (pip install -U jax), NVIDIA GPU (pip install -U "jax[cuda13]"), and Google TPU (pip install -U "jax[tpu]"). AMD and Intel GPU support is available with additional setup steps documented in the repository.