Before the first call

Chapter 03 introduces lower and compile as a way to see a program before running it. There are four stops on that road, not two, and running them separately tells you something the combined call cannot: which stage you are actually waiting on.

For a small MLP forward pass at 128 by 512 by 128, the answer on this machine is not close. Shape inference took 3.83 milliseconds, tracing 2.22, lowering to StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → 47.94, and XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → compiling that module 205.29. Everything upstream of the compiler is a rounding error against the compiler.

That ratio is what makes shape work cheap and worth doing eagerly. f.eval_shape(*specs) runs the abstract evaluation and hands back the output's shape and dtype without producing a program at all, which makes it the right tool for validating that a model's shapes line up across a config sweep. Nothing gets compiled, so nothing costs 205 milliseconds.

import time

import jax
import jax.numpy as jnp

spec = lambda *s: jax.ShapeDtypeStruct(s, jnp.float32)
args = (spec(128, 512), spec(512, 512), spec(512, 128))

def make(bias):
    def net(x, w1, w2):
        h = jnp.tanh(x @ w1 + bias)
        return jnp.sum(h @ w2)
    return jax.jit(net)

for warm in (0.0, 1.0):                    # warm every stage, then throw it away
    make(warm).trace(*args).lower().compile()

def ms(call):
    t0 = time.perf_counter()
    out = call()
    return round(1e3 * (time.perf_counter() - t0), 2), out

t_eval, shape = ms(lambda: make(2.0).eval_shape(*args))
f = make(3.0)
t_trace, traced = ms(lambda: f.trace(*args))
t_lower, lowered = ms(lambda: traced.lower())
t_compile, compiled = ms(lambda: lowered.compile())
t_again, _ = ms(lambda: lowered.compile())

print("eval_shape   ", t_eval, "ms ->", shape)
print("trace        ", t_trace, "ms ->", type(traced).__name__)
print("lower        ", t_lower, "ms ->", type(lowered).__name__)
print("compile      ", t_compile, "ms ->", type(compiled).__name__)
print("compile again", t_again, "ms, the same Lowered")

# eval_shape    3.83 ms -> ShapeDtypeStruct(shape=(), dtype=float32)
# trace         2.22 ms -> Traced
# lower         47.94 ms -> Lowered
# compile       205.29 ms -> Compiled
# compile again 0.02 ms, the same Lowered

Each stage hands back an object, and each object answers a different class of question. Traced carries the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, so it answers what JAX recorded. Lowered carries the StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → and can be printed with as_text, so it answers what the compiler was handed. Compiled carries the executable, so it answers what the compiler decided, in bytes and in flops.

One of those answers this arc has already used: the signature diff in lesson two came off a Lowered. The other two are read elsewhere. LAB·J2 takes its flop count off a Compiled, and so does the aliased-byte count that proves a buffer donation, in chapter 11's lesson on buffers you promise away.

The one caution worth carrying is the one the AOT page states about itself: these analysis results are diagnostics, not a stable interface, and their type and contents may differ across backends and versions. Treat a number from memory_analysis or cost_analysis as evidence about this build on this backend, not as a documented value you can assert against.

The last row of that table is the surprising one. Calling .compile() twice on the same Lowered object cost 205.29 milliseconds and then 0.02, so something below jit is caching compiled modules keyed on the module itself, not on the Python function that produced it.

That is a different cache from the one lesson one instrumented, and it explains a pattern that otherwise looks like magic: two unrelated functions that happen to lower to identical StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → pay for one compile between them. It also means a benchmark that compiles the same module twice and reports the second number is measuring the cache, not the compiler.

In a single process this cache is a convenience. Across processes it is configuration, since JAX also ships a persistent compilation cache that writes compiled modules to a directory and reads them back on the next run. That one is off unless you turn it on, and it is the difference between paying compile time once per run and once ever.

A Compiled is not a jitted function, and the difference shows up the moment you call it with something it was not compiled for. A jitted function meeting a new shape traces and compiles a new entry, silently; a Compiled meeting a new shape raises.

The message is specific about both sides of the mismatch, naming the argument, the type it was compiled with, and the type it was called with. That is the behaviour you want at a serving boundary, where a silent recompile at a request boundary is a latency spike nobody asked for.

Read that against lesson one and the arc closes. jit's cache is a convenience that hides a decision, and the decision is which of four things changed. The ahead-of-time surface removes the hiding: you choose the signature, you pay for the compile when you choose to, and anything that does not match gets an error instead of a recompile.

jit recompiles quietly. A compiled program refuses out loud.

import jax
import jax.numpy as jnp

f = jax.jit(lambda x: (x @ x).sum())
compiled = f.lower(jax.ShapeDtypeStruct((256, 256), jnp.float32)).compile()

print(compiled(jnp.ones((256, 256))))
for arg in (jnp.ones((128, 128)), jnp.ones((256, 256), jnp.bfloat16)):
    try:
        compiled(arg)
    except TypeError as e:
        print(f"{type(e).__name__}: {e}")

# 16777216.0
# TypeError: Argument types differ from the types for which this computation was compiled. The mismatches are:
# Argument 'x' compiled with float32[256,256] and called with float32[128,128]
# TypeError: Argument types differ from the types for which this computation was compiled. The mismatches are:
# Argument 'x' compiled with float32[256,256] and called with bfloat16[256,256]