TopolearnReading the intermediate program explains recompiles, unexpected ops and performance at the compiler level.
Every transformation works on an intermediate language called jaxpr. jax.make_jaxpr(f)(x) shows it: typed input variables, a list of primitive equations (mul, sin, dot_general, pjit, scan...) and outputs. Python control flow is gone; only primitives on abstract values remain. Reading jaxprs is the quickest way to see what tracing captured, e.g. whether a loop was unrolled or a constant was baked in.
Transformations are interpreters over jaxprs: grad linearizes and transposes them, vmap rewrites each primitive with a batching rule, jit lowers them.
Ahead-of-time stages make the pipeline explicit:
jax.jit(f).trace(x)gives the traced stage,.lower(x)lowers to StableHLO (MLIR);lowered.as_text()prints it..compile()runs XLA and returns aCompiledobject;compiled.as_text()shows the optimized HLO after fusion, andcompiled.cost_analysis()estimates flops and bytes.jax.exportserializes a lowered function so it can be loaded and run later or elsewhere.
Every primitive has rules for evaluation, abstract evaluation (shape inference), lowering, JVP, transpose and batching. That is why adding new low-level ops is done through jax.extend or custom calls (FFI), and why the custom-kernel language Pallas plugs in as a primitive too.
Captured constants matter: a large NumPy array closed over by a jitted function is embedded in the program, inflating compile time and memory. Pass it as an argument instead.
Example
import jax
import jax.numpy as jnp
def f(x, y):
return jnp.sin(x) * y + 1.0
print(jax.make_jaxpr(f)(1.0, 2.0))
# { lambda ; a:f32[] b:f32[]. let
# c:f32[] = sin a
# d:f32[] = mul c b
# e:f32[] = add d 1.0:f32[]
# in (e,) }
lowered = jax.jit(f).lower(1.0, 2.0)
print("stablehlo.sine" in lowered.as_text()) # True: StableHLO module
compiled = lowered.compile()
print(compiled(1.0, 2.0)) # 2.682942
print(jax.make_jaxpr(jax.grad(f))(1.0, 2.0).eqns[0].primitive) # sinRemember
- A jaxpr is the primitive-level program every transformation operates on
.lower()gives StableHLO;.compile()runs XLA and exposes optimized HLO and cost analysis- Closed-over large constants get baked into the program; pass them as arguments
Watch out for
- Closing over a big NumPy array in a jitted function, which embeds it as a constant in every compile.