JAX is a numerical computing library that applies four composable function transformations—jit (XLA compilation), grad (automatic differentiation), vmap (vectorization), and pmap (device parallelism)—to pure Python functions written against a NumPy-compatible API, producing programs that run at near-peak efficiency on CPUs, GPUs, and TPUs without requiring any framework-specific model class or imperative training loop.

<svg viewBox="0 0 760 470" xmlns="http://www.w3.org/2000/svg" font-family="ui-monospace,monospace">
  <rect width="760" height="470" fill="#0d1117"/>
  <text x="380" y="32" text-anchor="middle" fill="#e6edf3" font-size="21" font-weight="700">JAX Composable Transformations</text>
  <text x="380" y="52" text-anchor="middle" fill="#8b98a5" font-size="13">Four transforms stack on any pure function — jit∘grad∘vmap∘pmap all compose</text>

  <!-- Center function box -->
  <rect x="290" y="72" width="180" height="50" rx="6" fill="#161b22" stroke="#e6edf3" stroke-width="1.5"/>
  <text x="380" y="96" text-anchor="middle" fill="#e6edf3" font-size="13" font-weight="700">def f(params, x):</text>
  <text x="380" y="112" text-anchor="middle" fill="#8b98a5" font-size="11">pure Python + jnp ops</text>

  <!-- jit arrow -->
  <line x1="380" y1="122" x2="380" y2="158" stroke="#3fb950" stroke-width="1.5" stroke-dasharray="4,2"/>
  <rect x="290" y="160" width="180" height="44" rx="4" fill="#238636" fill-opacity="0.2" stroke="#238636" stroke-width="1.5"/>
  <text x="380" y="178" text-anchor="middle" fill="#3fb950" font-size="12" font-weight="700">jit(f)</text>
  <text x="380" y="196" text-anchor="middle" fill="#8b98a5" font-size="10">trace → XLA HLO → compile (~500 ms first call)</text>

  <!-- grad arrow (left) -->
  <line x1="290" y1="182" x2="180" y2="182" stroke="#58a6ff" stroke-width="1.5" stroke-dasharray="4,2"/>
  <rect x="30" y="160" width="148" height="44" rx="4" fill="#1f6feb" fill-opacity="0.2" stroke="#1f6feb" stroke-width="1.5"/>
  <text x="104" y="178" text-anchor="middle" fill="#58a6ff" font-size="12" font-weight="700">grad(f)</text>
  <text x="104" y="196" text-anchor="middle" fill="#8b98a5" font-size="10">reverse-mode AD · O(1) extra memory</text>

  <!-- vmap arrow (right) -->
  <line x1="470" y1="182" x2="582" y2="182" stroke="#f0883e" stroke-width="1.5" stroke-dasharray="4,2"/>
  <rect x="582" y="160" width="148" height="44" rx="4" fill="#f0883e" fill-opacity="0.2" stroke="#f0883e" stroke-width="1.5"/>
  <text x="656" y="178" text-anchor="middle" fill="#f0883e" font-size="12" font-weight="700">vmap(f)</text>
  <text x="656" y="196" text-anchor="middle" fill="#8b98a5" font-size="10">batched kernel · 0 Python loop overhead</text>

  <!-- Composition row -->
  <text x="380" y="236" text-anchor="middle" fill="#8b98a5" font-size="12">Composing all four — the standard research pattern</text>
  <rect x="80" y="248" width="600" height="44" rx="4" fill="#161b22" stroke="#a371f7" stroke-width="1.5"/>
  <text x="380" y="266" text-anchor="middle" fill="#a371f7" font-size="12" font-weight="700">jit(vmap(grad(f)))</text>
  <text x="380" y="284" text-anchor="middle" fill="#8b98a5" font-size="10">batched gradient · XLA-compiled · ~1 µs Python dispatch after first call</text>

  <!-- pmap row -->
  <rect x="80" y="304" width="600" height="44" rx="4" fill="#161b22" stroke="#f85149" stroke-width="1.5"/>
  <text x="380" y="322" text-anchor="middle" fill="#f85149" font-size="12" font-weight="700">pmap(jit(vmap(grad(f))))  — across 8 TPU/GPU devices</text>
  <text x="380" y="340" text-anchor="middle" fill="#8b98a5" font-size="10">SPMD data-parallel · all_reduce via XLA collective · 275 TFLOPS bfloat16 per TPU v4 chip</text>

  <!-- XLA compilation pipeline -->
  <text x="30" y="374" fill="#8b98a5" font-size="12">XLA compilation pipeline</text>
  <rect x="30"  y="382" width="100" height="28" rx="3" fill="#161b22" stroke="#30363d" stroke-width="1"/>
  <rect x="140" y="382" width="100" height="28" rx="3" fill="#161b22" stroke="#30363d" stroke-width="1"/>
  <rect x="250" y="382" width="100" height="28" rx="3" fill="#161b22" stroke="#30363d" stroke-width="1"/>
  <rect x="360" y="382" width="100" height="28" rx="3" fill="#161b22" stroke="#30363d" stroke-width="1"/>
  <rect x="470" y="382" width="100" height="28" rx="3" fill="#161b22" stroke="#30363d" stroke-width="1"/>
  <text x="80"  y="400" text-anchor="middle" fill="#e6edf3" font-size="10">Python trace</text>
  <text x="190" y="400" text-anchor="middle" fill="#e6edf3" font-size="10">HLO graph</text>
  <text x="300" y="400" text-anchor="middle" fill="#e6edf3" font-size="10">HLO optimize</text>
  <text x="410" y="400" text-anchor="middle" fill="#e6edf3" font-size="10">LLVM / PTX</text>
  <text x="520" y="400" text-anchor="middle" fill="#e6edf3" font-size="10">kernel cache</text>
  <line x1="130" y1="396" x2="140" y2="396" stroke="#8b98a5" stroke-width="1" marker-end="url(#a2)"/>
  <line x1="240" y1="396" x2="250" y2="396" stroke="#8b98a5" stroke-width="1" marker-end="url(#a2)"/>
  <line x1="350" y1="396" x2="360" y2="396" stroke="#8b98a5" stroke-width="1" marker-end="url(#a2)"/>
  <line x1="460" y1="396" x2="470" y2="396" stroke="#8b98a5" stroke-width="1" marker-end="url(#a2)"/>
  <defs><marker id="a2" markerWidth="4" markerHeight="4" refX="2" refY="2" orient="auto"><path d="M0,0 L4,2 L0,4 Z" fill="#8b98a5"/></marker></defs>

  <text x="380" y="460" text-anchor="middle" fill="#6e7681" font-size="11">Shape change → recompile; same shape → ~1 µs cached dispatch · Flax, Optax, Equinox build on these four transforms</text>
</svg>

JAX's defining architectural choice is that every transformation—differentiation, vectorization, parallelism, compilation—is a higher-order function that takes a pure Python function and returns a new pure Python function, making them composable in any order without any framework registry, decorator magic, or special model class. jit(vmap(grad(f))) is legal JAX: grad wraps f and returns its gradient function, vmap wraps that gradient to run it over a batch dimension simultaneously, and jit traces the batched gradient through XLA and compiles it to a single fused GPU/TPU kernel. The composition rule requires only that each inner function be pure—no Python-level side effects, no mutation of external state—a constraint JAX enforces via abstract value tracing rather than runtime checks.

The XLA compiler that backs jit converts a traced JAX function into a High-Level Operations (HLO) graph, applies algebraic simplifications and fusion, then emits LLVM IR for CPU or PTX for GPU—a process that takes approximately 500 ms for a simple function on first call and zero compile time on every subsequent call with the same argument shapes. After compilation, Python dispatch overhead is approximately 1 µs, so a JAX training step that takes 50 ms on GPU spends only 0.002% of its time in Python. The compilation result is keyed on both the function and the argument shape tuple: calling a compiled function with a differently-shaped input triggers recompilation. For models whose sequence length varies per batch, jax.jit(f, static_argnums=()) with explicit padding to a fixed length avoids per-step recompilation.

Automatic differentiation via jax.grad implements reverse-mode backpropagation by recording a trace of primitive operations during the forward pass and evaluating their vector-Jacobian products in reverse, at a memory cost that is O(1) extra storage relative to the forward pass rather than O(N layers) as in eager frameworks. A 1-billion-parameter model gradient computed via grad requires approximately 4 GB in bfloat16. Composing grad twice yields the second derivative; jax.hessian(f) builds the full N×N Hessian by applying jacfwd(jacrev(f)), with memory cost O(N²) for N parameters. For large models where storing all activations is prohibitive, jax.checkpoint (gradient checkpointing) reduces peak activation memory from O(N) to O(√N) for a 100-layer model—from 100 units to 10 units—at the cost of recomputing activations during the backward pass, adding approximately 33% more FLOPs. A ResNet-50 gradient without checkpointing requires ~800 MB of activation memory; with checkpointing this falls to ~90 MB.

Vectorization via vmap eliminates the per-sample Python loop that plagues research code by automatically vectorizing a scalar function over a leading batch dimension, producing a single batched XLA kernel with zero Python loop overhead and exactly the same asymptotic performance as a manually batched implementation. The canonical use is jax.vmap(jax.grad(loss))(params, x_batch): computing per-sample gradients, which in PyTorch requires either an explicit loop or functorch.vmap, becomes a single line in JAX. vmap can also be nested—vmap(vmap(f)) vectorizes over two independent batch axes simultaneously—and combined with grad to produce per-sample Jacobians for influence function computation and natural gradient methods.

Physical device parallelism via pmap maps a function across multiple accelerators using SPMD (Single Program, Multiple Data) parallelism, where each device executes identical code on a shard of the data and collective operations like jax.lax.pmean synchronize results across devices via XLA's all-reduce. A TPU v4 chip delivers 275 TFLOPS in bfloat16 and 137 TFLOPS in float32; a TPU v4 pod with 4,096 chips reaches approximately 1.13 EFLOPS—numbers that JAX's pmap can saturate because pmap compiles the full multi-device program as a single XLA computation rather than coordinating devices through a Python scheduler. The bfloat16 format (1 sign, 8 exponent, 7 mantissa bits) matches float32 in dynamic range while halving memory bandwidth—a 7B-parameter model requires ~14 GB in bfloat16 versus ~28 GB in float32—making it the preferred dtype for large-model training where memory bandwidth limits throughput more than compute precision.

Sequence loops in JAX must use lax.scan rather than Python for loops with a functional fold that compiles in O(1) time regardless of the number of iterations, and is the correct primitive for any recurrent computation in JAX. A Python for loop inside jit is unrolled at trace time: a 1,000-iteration loop produces an HLO graph with 1,000 copies of the loop body, inflating compile time from ~500 ms to ~500 s; a 10,000-iteration loop raises compile time to over 5,000 s. lax.scan(f, init, xs) compiles the body function once and instructs XLA to iterate it—compile time stays constant at ~500 ms and memory usage stays O(1) for the carry state regardless of sequence length. An LSTM over 10,000 timesteps compiles in ~500 ms and runs in ~30 ms on GPU, versus ~5,000 s compile for an unrolled loop. The same principle applies to lax.cond (replaces Python if on dynamic values) and lax.while_loop (variable-length iteration with a compiled loop predicate).

JAX EXECUTION FLOWCHART

  Python def f(params, x): return jnp.dot(params, x)
         │
         ▼
  ┌─────────────────────┐
  │  jit(grad(vmap(f))) │  transform application (O(1) Python cost)
  │  = new Python fn    │  no compilation yet
  └────────┬────────────┘
           │  first call with concrete shapes
           ▼
  ┌─────────────────────┐
  │  Abstract tracing   │  ShapedArray values (no actual data)
  │  → HLO graph        │  ~500 ms compile (CPU→LLVM, GPU→PTX)
  └────────┬────────────┘
           │  compiled kernel cached by (fn_id, shapes)
           ▼
  ┌─────────────────────┐
  │  Subsequent calls   │  ~1 µs Python dispatch
  │  same shapes → hit  │  kernel runs on XLA device
  └────────┬────────────┘
           │  shape change?
           ▼
  ┌─────────────────────┐
  │  Recompile          │  new shape → new HLO → new cache entry
  │  (avoid with pad)   │  pad sequences to fixed length
  └─────────────────────┘

Read JAX through a functional transformation lens rather than a deep learning framework lens. PyTorch and TensorFlow provide modules, optimizers, and data loaders as first-class objects; JAX provides four composable mathematical transformations on pure functions, and leaves the model, optimizer, and training loop as ordinary Python code. Flax, Optax, and Equinox are not JAX extensions—they are Python libraries that happen to work well with functions, pytrees, and the four transforms, and any code that satisfies the purity constraint can use them. That design decision is what makes JAX uniquely suited to research that pushes outside the standard supervised-learning loop: meta-learning, neural ODEs, physics-informed networks, and hardware-aware kernel design all benefit from treating gradients, vectorization, and device placement as first-class composable operations rather than as framework features.