Most numerical Python gives you two bad choices. Either you write eager code that is easy to debug but slow, or you write a static graph that is fast but alien. JAX’s answer is a third road: keep writing NumPy-style Python, but capture it — trace it — into a small, serializable intermediate representation the moment you ask for speed. That capture step is why the same grad, jit and vmap decorators work unchanged from a laptop CPU to a TPU pod. Where PyTorch’s runtime, which we covered in the PyTorch source tour, records ops while they execute, JAX re-executes your Python function with tracer objects and records a pure program instead.

In this post we walk the google/jax source tree (version 0.12.1 at the time of writing) along one axis only: what happens between @jax.jit on your function and device buffers coming back. That path runs through four neighborhoods — the public API, the tracing machinery that builds a jaxpr, the MLIR/XLA lowering and compilation stage, and the sharded execution runtime. The C++ half of the project, jaxlib, lives mostly outside the Python tree we will link here, but you will see its fingerprints in jax._src.lib, the extension module that wraps the XLA client.

If you have read our pieces on ONNX graph interchange or on tinygrad, the jaxpr will feel familiar: it is a flat, SSA-style list of equations, much closer to an ONNX graph than to Python. The difference is that JAX never asks you to write that graph. The Python source is the front end, and everything below it is infrastructure you can read like any other program.

Overview of the JAX source, from the public API through tracing, lowering and compilation down to sharded execution. Generated from the repository tree of google/jax on branch main.

Reading the overview from left to right:

• Everything starts at jit, L187 in the public API module — and despite the name, it immediately hands off to make_jit, L453 in the pjit module. Modern JAX has exactly one implementation of “compile this function”, and jit is a thin shell over it.
• Argument handling happens twice: once as Python-level flattening in flatten_fun, L72, and once as pytree decomposition in tree_flatten, L86. JAX sees your dict of parameters as a tree of leaves plus a treedef, and that structure survives all the way into the compiled artifact.
• The transformation stack lives in WrappedFun, L145: every transform (grad, vmap, remat) wraps the function with parameters that the tracer consults later, so transforms compose by nesting wrappers rather than by editing code.
• Tracing itself is JaxprTrace, L116: a Python interpreter that follows your function while collecting JaxprEqn, L411 nodes into a Jaxpr, L99.
• Autodiff is not a second system — linearize, L227 re-traces your function under a JVP tracer and reuses the same partial-evaluation machinery, which is why grad and jit compose in any order.
• Lowering turns the jaxpr into an MLIR module inside StableHLO dialects at lower_jaxpr_to_module, L1394, using per-primitive rules registered by the lax library.
• Compilation is compile_or_get_cached, L436: it resolves the platform client through get_backend, L980, checks the persistent cache, and otherwise ships the module to XLA.
• Execution has two doors: the eager op-by-op path in apply_primitive, L82 (which, amusingly, works by jitting each single primitive), and the sharded path, where shard_args, L86 slices inputs across the mesh before a MeshExecutable runs.
• What comes back is an ArrayImpl, L226: a NumPy-like object that is really a bag of device buffers, each with its own sharding.

Tracing: Where Python Becomes a Graph

The core trick deserves a closer look, because everything else in JAX is built on it. A tracer is a stand-in value that records every operation applied to it. When you call a jitted function, JAX wraps your arguments in tracers, runs your Python code, and harvests the record. The result is the jaxpr: invariants in, constants folded in, and a straight-line list of equations, each naming a primitive, its inputs and its outputs. You can see one yourself without any ceremony:

import jax
import jax.numpy as jnp

def f(x):
    return jnp.sin(x) * 2.0

print(jax.make_jaxpr(f)(jnp.zeros((3,))))

That output — make_jaxpr is defined at make_jaxpr, L2194 — is a real data structure you can pass around, transform and inspect. Control flow like lax.scan and lax.while_loop shows up as higher-order equations holding nested jaxprs, which is precisely why Python for loops over data do not survive tracing: the tracer only sees unrolled operations unless you reach for those primitives. Each equation’s meaning is defined twice, once abstractly and once concretely; the concrete evaluator is eval_jaxpr, L748, which is what runs when you ask for eager execution.

Primitives are declared in one place and given rules everywhere else, and the lax library is the registry. Look at how a basic op is assembled at add_p, L5524: a primitive object, then three lines that attach its JVP rule to ad.primitive_jvps, its transposition rule to ad.primitive_transposes, and its XLA lowering via mlir.register_lowering. Every transform you use is a dictionary lookup against tables like these. This separation — one primitive, many small rules — is the quiet design decision that makes the whole system composable, and it is a pattern worth stealing for your own DSLs. Source provenance is not lost either: equations carry file/line data through SourceInfo, L136, so stack traces through three nested transforms still point at your original line of Python.

Autodiff and the Transform Stack

Because grad is “run the tracer, then differentiate the jaxpr”, JAX autodiff is source-to-source in spirit even though it never touches your Python text. linearize, L227 splits a function into its primal pass and a linear part; transposition turns the linear part around to get vector-Jacobian products. Reverse-mode differentiation of a jaxpr is effectively a clever backwards interpreter over the same equation list, which is why memory, not speed, is the binding constraint — the whole primal trace stays live until the backward pass consumes it.

That constraint has an elegant fix: jax.checkpoint (remat) is defined at remat, L393 and works by re-tracing the forward function inside the backward pass, trading recomputation for memory. It composes with everything because it is implemented through the same partial-evaluation hooks as any other transform. The same story repeats for vmap: BatchTrace, L232 adds a batch dimension to every operation via per-primitive batching rules, and custom rules hook in through the same primitive tables. Even user-defined derivatives, declared with custom_jvp, L87, integrate by registering rules rather than by special-casing.

If you have trained models with transformers or ported a nanogpt loop, this is the machinery you were using: loss = jax.value_and_grad(loss_fn) is two nested tracers deep before it ever touches XLA. And when you stack jit(grad(vmap(f))), you are literally nesting WrappedFun wrappers — the design keeps each layer dumb and the composition correct.

From Jaxpr to Silicon: Lowering and Compilation

The jaxpr leaves Python-land at the MLIR boundary. lower_jaxpr_to_module walks the equations and calls each primitive’s registered lowering rule, which typically emits StableHLO operations; constants become embedded literals via ir_constants, L285 machinery around it. You can drop from this level to hand-written kernels when needed — mixing in Triton via jax.experimental.pallas or shaping how XLA fuses things — but the default path assumes XLA will do the fusion, layout and scheduling work that a hand-tuned stack like CUTLASS does manually.

Compilation proper is where the C++ world takes over. compile_or_get_cached, L436 gathers compile options at get_compile_options, L187, consults the persistent cache — keys built in VerificationCache, L131 — and finally calls backend_compile_and_load against a client obtained from the backend registry. That registry, xla_bridge, maps a platform string ("cpu", "gpu", "tpu") to a lazily-initialized XLA client; get_backend at get_backend, L980 is the doorway. In architecture terms this is the same boundary an ONNX Runtime execution provider or an IREE target backend occupies — compare our onnxruntime and IREE tours — except JAX’s “graph format” is the jaxpr-to-StableHLO lowering rather than a serialized file.

One detail worth appreciating: compilation is keyed on shapes, dtypes, sharding and static arguments, not on values. That is why the first call to a jitted function is slow and every structurally identical call after it is free, and why changing a Python int that was marked static_argnums silently triggers a recompile.

From Install to Your First Model

The CPU wheel is a two-liner, and the CUDA one is only slightly longer:

pip install -U jax            # CPU only
pip install -U "jax[cuda12]"  # NVIDIA GPU plugin wheels

Everything shown below uses only stable public API — the same calls that appear in the library’s own docstrings:

import jax
import jax.numpy as jnp

key = jax.random.key(0)                      # explicit PRNG state
params = {"w": jax.random.normal(key, (3,)),
          "b": jnp.zeros(())}

def loss(params, x, y):
    pred = x @ params["w"] + params["b"]
    return jnp.mean((pred - y) ** 2)

x = jax.random.normal(key, (64, 3))
y = jnp.ones((64,))

grad_fn = jax.jit(jax.grad(loss))            # grad composed with jit
g = grad_fn(params, x, y)
step = jax.tree.map(lambda p, e: p - 0.1 * e, params, g)
print(jax.jit(loss)(params, x, y))

Three things to notice. First, jax.random.key(0) returns an explicit key array — the PRNG type defined at PRNGKeyArray, L142, backed by the counter-based threefry_2x32, L243 bit generator. Randomness is a first-class, splittable value instead of a global generator, which is what makes seeded, reproducible training across devices tractable. Second, grad handled the dictionary params transparently because gradients flow through pytrees. Third, the update used jax.tree.map rather than a loop: under jit, ordinary Python iteration over data would be unrolled into the trace. From here, adding a real model is just more jnp — and if you want to see where such a loop goes when it grows up, our inference and vLLM pieces cover the serving side of the same ecosystem.

Sharding: One Function, Many Devices

JAX’s scaling story is not a separate API bolted on the side; it is metadata on the same jaxpr. You describe a logical Mesh, L235 of devices, annotate arrays with a NamedSharding, L71 that names mesh axes, and let the compiler insert the collectives. pjit — the engine underneath jit — traces your function, propagates sharding constraints through the jaxpr, and compiles one program whose inputs and outputs may live split across the whole mesh. Inside the runtime, shard_args, L86 handles slicing host-side inputs into per-device buffers, and the executable that comes back is a MeshExecutable, L1846 that fans the launch out across devices.

For per-shard programming, shard_map, L106 drops you one level down: your function runs on each slice with explicit collective primitives for communication. This is the mental model behind large-scale training and serving on TPUs — the same “partition the tensors, minimize the chatter” problem that ZeRO-style optimizers such as DeepSpeed solve for the PyTorch world, but expressed as compiler-visible annotations instead of runtime wrappers. The payoff is that resharding between two operations is inferred by the compiler rather than hand-written, and an ArrayImpl, L226 hides whether it is replicated, split, or scattered across an entire pod.

Honest Limits

The tracing model is powerful, and it is also the source of every common complaint. Tracing only sees values’ shapes and dtypes, so data-dependent Python control flow (if x > 0: on an array, for over data) either fails or unrolls; you must express real control flow with lax.while_loop and lax.scan, which is a genuine change of habit. Every new shape is a new compilation, so workloads with dynamic sequence lengths pay recompile taxes unless you pad or bucket — JIT-friendly design is a skill, not a default. The first call latency of a large jitted function can be seconds, since the whole lowering-to-XLA path runs synchronously, and error messages arriving from deep inside compiled code are much better than they used to be but still speak StableHLO dialect rather than Python. Eager mode exists, but it is best treated as a debugging tool, not a production mode. Memory predictability cuts both ways: because reverse-mode keeps the primal trace alive, OOMs arrive unless you plan remat and donation deliberately. And while the public NumPy surface is famously stable, the internals we toured — pxla, pjit internals, the lowering cache — move constantly; link against them at your peril. None of these are disqualifying, but they explain why JAX feels different in the hand from eager frameworks: you are not calling an operator library, you are writing programs for a compiler, and the source tree we just walked makes that explicit at every layer.

Enjoyed this post? Never miss out on future posts by following us