# JAX

JAX is Google's library for NumPy-style computing with automatic differentiation, JIT compilation via XLA, vectorisation and parallelism. Learn how it works.

## Overview

JAX is an open-source Python library for fast numerical computing on CPUs, GPUs and TPUs. Its array API, `jax.numpy`, looks almost exactly like NumPy. What makes it different is a set of function transformations. `grad` turns a function into one that returns its derivative. `jit` compiles a function with the XLA compiler. `vmap` makes a function that works on one example work on a whole batch. Sharding tools spread a computation across many accelerators.

These transformations compose. You can take the gradient of a vectorised function, compile it and run it across hundreds of chips, each step being one more wrapper around the same plain Python function. JAX itself does not provide neural network layers; libraries such as Flax and Optax are built on top of it for that.

An everyday picture is a set of power-tool attachments that all fit the same drill. One attachment sands, one drills, one polishes. You can fit several in sequence. JAX's transformations are attachments for functions: make it differentiable, make it batched, make it compiled, make it parallel.

## What it is

JAX's core is a small set of transformations plus an array library.

| Piece | What it does |
|---|---|
| `jax.numpy` | NumPy-compatible arrays and functions that run on accelerators |
| `jax.grad` | Returns a function that computes gradients of a scalar function |
| `jax.value_and_grad` | Returns both the value and the gradient in one pass |
| `jax.jit` | Compiles a function with XLA for speed |
| `jax.vmap` | Automatically vectorises a function over a batch axis |
| Sharding and `jax.sharding.Mesh` | Describe how arrays are split across devices so `jit` can parallelise |
| `jax.shard_map` | Write a per-device program with explicit communication |
| `jax.random` | Random numbers driven by explicit keys |
| `jax.lax` | Lower-level primitives, including structured loops and conditionals |

Around JAX sits an ecosystem of separate projects: Flax for neural network modules (its newer NNX API is recommended for new users), Optax for optimizers and gradient transformations, Orbax for checkpointing and Grain for data loading. [Keras](/tools/keras) can also use JAX as a backend.

## Where it came from

JAX builds on Autograd, a library introduced in 2015 by Dougal Maclaurin, David Duvenaud and Ryan P. Adams that could differentiate ordinary Python and NumPy code. At Google, that idea was carried forward and combined with XLA, the compiler Google had built for TensorFlow and TPUs.

At SysML 2018, Roy Frostig, Matthew James Johnson and Chris Leary of Google Brain presented the early system in a paper titled "Compiling machine learning programs via high-level tracing". JAX was released as open source under the Apache 2.0 licence.

Adoption grew mainly in research. DeepMind described wide internal use of JAX in December 2020. The PaLM paper of 2022 reported training a 540-billion-parameter language model with a JAX-based system on 6,144 TPU v4 chips. Parallelism APIs evolved over time: the early `pmap` was rebuilt on top of `jit` and `shard_map` in version 0.8.0 in October 2025 and placed in maintenance mode, with `jax.jit(jax.shard_map(...))` recommended for new code. JAX still uses 0.x version numbers and makes regular releases.

## How it works

### Step 1: tracing

When you call a transformed function, JAX does not simply run the Python. It passes in special tracer objects that record every array operation. The result is an intermediate representation called a jaxpr, a simple list of primitive operations with their inputs and outputs.

### Step 2: transforming

Each transformation rewrites that representation. `grad` adds the reverse-mode derivative operations, which is [backpropagation](/tools/backpropagation) done mechanically. `vmap` rewrites operations to carry an extra batch dimension. Because each transformation takes a function and returns a function, they can be stacked in any order.

### Step 3: compiling

`jit` hands the representation to XLA, which fuses operations, plans memory and generates code for the target hardware. The first call compiles; later calls with the same input shapes and types reuse the compiled program, which is why shapes should stay stable.

### Step 4: parallelising

For many devices you define a mesh, a named grid of devices, and annotate how arrays are sharded across it. `jit` then partitions the computation and inserts the communication automatically. If you need exact control, `shard_map` lets you write the program each device runs and call collectives yourself. These map onto ideas such as [data parallelism](/tools/data-parallelism) and [tensor parallelism](/tools/tensor-parallelism).

### Pure functions and explicit state

Tracing only works reliably for pure functions: the output must depend only on the inputs, with no hidden side effects. So JAX makes state explicit. Model parameters are passed in as arguments and new parameters are returned. Random numbers come from keys you split and pass around.

```python
import jax, jax.numpy as jnp

def loss(w, x, y):
    return jnp.mean((x @ w - y) ** 2)

grad_fn = jax.jit(jax.grad(loss))
w = jnp.zeros(3)
x, y = jnp.ones((8, 3)), jnp.ones(8)
w = w - 0.1 * grad_fn(w, x, y)
```

## Key concepts and terms

| Term | Meaning |
|---|---|
| Transformation | A function that takes a function and returns a modified one |
| `grad` | Transformation that produces a gradient function |
| `jit` | Just-in-time compilation with XLA |
| `vmap` | Automatic vectorisation over a batch axis |
| XLA | Accelerated Linear Algebra, the compiler JAX uses |
| Tracer | A stand-in object that records operations during tracing |
| jaxpr | JAX's intermediate representation of a traced function |
| Pure function | A function with no side effects whose output depends only on inputs |
| PRNG key | Explicit state for random number generation |
| Pytree | A nested structure of lists, dicts and tuples that JAX treats as a container of arrays |
| Mesh | A named arrangement of devices used for sharding |
| `shard_map` | Per-device programming with explicit collectives |

## A worked example

This scenario is illustrative. A research group wants to fit a physical model with 50 parameters to 10,000 separate experiments, each with its own data, and then train a small neural network on the results.

1. They write a function `simulate(params, inputs)` for a single experiment using `jax.numpy`.
2. They define `loss(params, inputs, target)` for one experiment and wrap it with `jax.grad` to get gradients.
3. Instead of writing a loop over experiments, they apply `jax.vmap` so the gradient function handles all 10,000 experiments in one batched call.
4. They wrap the batched update in `jax.jit`. The first call takes a few seconds to compile; each later step runs as one fused program on the GPU.
5. They use Optax's Adam optimizer, which keeps its state in a pytree that they pass in and return each step.
6. For the neural network stage they move to a machine with eight GPUs, define a one-dimensional mesh and shard the batch across it. The model code stays the same; only the sharding annotations change.

The same small set of transformations covers batching, differentiation, compilation and scaling.

## Strengths and limitations

**Strengths:**

- Composable transformations make complex numerical code short.
- XLA compilation gives strong performance on GPUs and TPUs.
- Excellent scaling across large device meshes.
- Familiar NumPy-style API.
- Explicit state makes randomness and reproducibility clear.

**Limitations:**

- The pure-function style and explicit state take time to learn.
- Python control flow that depends on array values needs `jax.lax` constructs inside `jit`.
- Changing input shapes trigger recompilation.
- Fewer ready-made models and tutorials than [PyTorch](/tools/pytorch).
- Neural network features live in separate libraries whose APIs have changed over time.

## How it compares with alternatives

| Tool | Programming style | Compilation | Typical use |
|---|---|---|---|
| JAX | Functional transformations | XLA via `jit` | Research at scale, scientific computing, TPUs |
| [PyTorch](/tools/pytorch) | Imperative, object-oriented | Optional `torch.compile` | Most research and production deep learning |
| [TensorFlow](/tools/tensorflow) | Eager plus traced graphs | `tf.function`, XLA | Production pipelines and deployment |
| [Keras](/tools/keras) | High-level API | Uses the backend's compiler | Fast development, can run on JAX |
| NumPy | Imperative arrays | None, CPU only | General numerical work without gradients |

## Practical guidance and best practices

1. Write pure functions: take parameters and data as inputs and return new values.
2. Put `jit` around whole training steps, not around tiny functions.
3. Keep array shapes fixed across calls to avoid recompilation; pad batches if needed.
4. Split PRNG keys for every random operation and never reuse a key by accident.
5. Use `jax.lax.scan`, `cond` and `fori_loop` for loops and branches inside compiled code.
6. Use Flax for model structure and Optax for optimizers rather than writing them by hand.
7. For multiple devices, start with sharded `jit`; reach for `shard_map` only when you need explicit control.
8. Call `block_until_ready()` when timing code, because JAX dispatches work asynchronously.

## Common mistakes

- Using Python `print` inside a jitted function and seeing it run only once, at trace time.
- Mutating arrays in place; JAX arrays are immutable, so use `.at[].set()` instead.
- Reusing the same random key and getting identical "random" numbers.
- Branching with Python `if` on traced values, which raises errors under `jit`.
- Benchmarking the first call, which includes compilation time.
- Writing new code with `pmap` instead of the recommended sharding APIs.

## Frequently asked questions

### Is JAX a deep learning framework?

Not on its own. JAX is a numerical computing library with transformations. Deep learning frameworks such as Flax are built on it, and Keras can use it as a backend.

### Why does JAX require pure functions?

JAX traces functions once and then transforms and compiles the trace. Side effects such as printing or changing global variables are not part of the trace, so they behave unexpectedly. Pure functions make tracing correct and repeatable.

### Does JAX only work on TPUs?

No. JAX runs on CPUs, NVIDIA and AMD GPUs and TPUs. It is especially well known on TPUs because XLA was designed with them in mind.

### What replaced pmap?

Since version 0.8.0, `pmap` is built on `jit` and `shard_map` and is in maintenance mode. The documentation recommends `jax.jit` with sharded arrays, or `jax.jit(jax.shard_map(...))` when you need per-device code.

### How does randomness work in JAX?

You create a key with `jax.random.key`, split it into subkeys with `jax.random.split`, and pass a subkey to each random function. The same key always gives the same numbers, which makes experiments reproducible.

### Should I choose JAX or PyTorch?

Choose JAX if you value composable transformations, work at large scale on TPUs or write scientific code that needs gradients and batching. Choose PyTorch if you want the largest ecosystem of models and tools and a more familiar imperative style.

## Glossary

- Accelerator: hardware such as a GPU or TPU built for parallel arithmetic.
- Collective: a communication step across devices, such as an all-reduce.
- Flax: a neural network library built on JAX.
- Optax: a JAX library of optimizers and gradient transformations.
- Recompilation: building a new compiled program because input shapes or types changed.
- Reverse mode: the way of computing derivatives that backpropagation uses.
- Sharding: splitting an array across several devices.
- Vectorisation: running one operation over many data elements at once.

## Summary

JAX is a Google-born library that pairs a NumPy-like API with composable transformations for differentiation, compilation, vectorisation and parallelism. It grew from Autograd and XLA, was described at SysML 2018 and has powered large-scale research, including training runs across thousands of TPU chips. It rewards a pure-function style with speed and clean scaling, and it relies on libraries such as Flax and Optax for neural network building blocks.