JAX

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.

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 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 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 and 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.

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

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.
• Neural network features live in separate libraries whose APIs have changed over time.

How it compares with alternatives

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.

Related

It belongs to Frameworks in my stack, under Deep learning. Others in the same group: