JAX

What JAX is

JAX describes itself as a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.^[1] Both halves of that description matter. Array computation means the basic objects a program works on are multidimensional arrays — rectangular grids of numbers — and the operations on them are the familiar ones: elementwise arithmetic, matrix products, reductions along an axis. Program transformation means the library can accept an ordinary Python function and produce a modified version of it — compiled, batched, differentiated or parallelized — instead of asking the user to rewrite that function by hand.

JAX describes itself as narrowly scoped, focused on efficient array operations and program transformations.^[1] It does not present itself as a complete deep learning framework with ready-made layers, optimizers and training loops. The narrowness is deliberate: the machine learning pieces are supplied by libraries built on top, while JAX concentrates on making the underlying array machinery and transformations fast and predictable.

Transforming functions

JAX includes composable function transformations for compilation, batching, automatic differentiation and parallelization.^[1] In plain terms these do four different jobs. Compilation turns a Python function into code the accelerator can run as a unit, so that a long chain of small array operations does not pass through the Python interpreter at every step. Batching lets a function written for a single example be run over a whole stack of examples without an explicit loop. Automatic differentiation produces a function that returns the derivatives of the original one, which is what training a model requires. Parallelization spreads the work across several devices.

Composable means these transformations can be nested: a batched function can then be compiled, a compiled function can then be differentiated, and so on, in whatever order the problem needs. Writing such a combination by hand normally takes a large amount of specialized code; the point of the design is that it becomes a matter of wrapping one function in another.

One program, several backends

The same code executes on multiple backends, including CPU, GPU and TPU.^[1] That is a different arrangement from the array library most Python users know. NumPy operations run eagerly, synchronously and only on CPU, while JAX operations can run eagerly or after jit compilation, dispatch asynchronously and execute on CPU, GPU or TPU.^[2] Eager execution means each operation is carried out when it is called; with jit compilation the function is traced once and later calls reuse the compiled version. Asynchronous dispatch means the call returns before the device has finished, so the host program can keep going and only waits when it actually needs the result.

Tools for working at scale

JAX ships tools for ahead-of-time compilation, sharding and automatic parallelization, and per-device programming with shard_map.^[1] Ahead-of-time compilation means the compiled artifact is produced before the program starts running, rather than on first call. Sharding means splitting an array across devices, and automatic parallelization lets the compiler work out the communication that this split requires. shard_map offers a way to write code per device explicitly, for the cases where the automatic arrangement is not what the problem needs.

Compared with NumPy

The relationship to NumPy is the question users raise most often, and JAX's own documentation treats it as a performance question rather than a matter of principle. Microbenchmarks of single array operations on CPU generally favour NumPy because of its lower per-call dispatch overhead, while GPU or TPU workloads and JIT-compiled sequences generally favour JAX.^[2] The reason is scale: for one small operation on one array, the cost of dispatching the call dominates and NumPy's path is shorter. For long chains of operations, or for work on accelerators, compilation removes most of that overhead and the hardware does more at once.

The ecosystem and the learning path

Because JAX leaves model building to others, neural network libraries built on JAX include Flax, Equinox and Keras.^[1] The documentation organises learning into six levels, from JAX 101 on arrays and jax.grad to JAX 601 on the jaxpr language and Autodidax.^[1] The path starts with the parts most people use every day and ends with the lower-level machinery inside the transformation system.

Taken together, these choices give JAX a clear place: it is not a one-stop framework but a substrate that other libraries build on. What falls inside its scope — array operations and composable transformations — is made as fast and as uniform as possible; what falls outside it is left to the ecosystem.