← AI Terminology
JAX
JAX is a Google open-source numerical computing library that combines NumPy-compatible array operations with automatic differentiation (
It is the foundation of Google DeepMind's research stack and frameworks like Flax and Haiku.
grad), JIT compilation to XLA, and functional transformations (vmap, pmap) — designed for high-performance ML research on CPUs, GPUs, and TPUs.It is the foundation of Google DeepMind's research stack and frameworks like Flax and Haiku.
Why It Matters in AI
JAX's composable transformations make research code concise and hardware-agnostic:
jit(grad(vmap(f))) compiles, differentiates, and vectorises a function in one line. Its XLA backend achieves near-optimal performance on TPUs — giving DeepMind and Google Brain a competitive edge for large-scale research. Gemini models were trained primarily in JAX. It is increasingly popular for research where PyTorch's imperative style is too slow or inflexible.Key Points
| Aspect | Description |
|---|---|
| Backend | XLA (Accelerated Linear Algebra) compiler — same backend as TensorFlow TPU training |
| Frameworks | Flax (Google), Haiku (DeepMind), Optax (optimisers), Equinox — all built on JAX |
| vs PyTorch | More functional/research-oriented; steeper learning curve; PyTorch still dominates industry |
| TPU support | Best-in-class TPU performance — Google's primary framework for TPU pod training |
| Pure functions | JAX requires pure, stateless functions — no in-place mutation; state managed via pytree carry |
| Core transforms | jit (JIT compile), grad (autodiff), vmap (vectorise), pmap (parallelise across devices) |
Simple Analogy
NumPy on steroids: the same array math you know, but with a compiler that fuses operations into fast XLA kernels, an automatic differentiation engine, and a vectorisation pass — all composable like LEGO bricks.
Common Usage Examples
import jax.numpy as jnp; jnp.dot(a, b)— drop-in NumPy replacement with GPU/TPU accelerationgrad_fn = jax.grad(loss); grads = grad_fn(params, batch)— automatic differentiationjit_train_step = jax.jit(train_step)— compile training step to XLA for faster executionjax.vmap(predict)(batch_inputs)— vectorise prediction over a batch automatically- Flax model training:
state = TrainState.create(apply_fn=model.apply, params=params, tx=optax.adam(1e-3))
Summary
In short: JAX is Google's high-performance ML research framework — NumPy-compatible, composably transformable, and XLA-compiled — powering DeepMind research and Gemini model training at scale.