When JAX transforms a Python function, it first traces the computation into Jaxpr, a small, explicitly typed functional intermediate representation. Later transformations operate on this representation before compilation proceeds toward XLA. You can inspect it with jax.make_jaxpr(). section 5.2.1 section 5.1.2
A printed Jaxpr makes the computation’s structure visible:
- A lambda header separates constants from input variables; closed-over values appear among the constants. section 5.2.1
- Equations bind intermediate variables by applying primitives such as add, mul, sin, and reduce_sum. section 5.2.1
- Variables carry types and shapes, while outputs appear as atomic expressions. section 5.2.1
- Structured control flow can appear as higher-order primitives containing sub-Jaxprs, including cond, while, fori_loop, and scan. section 5.2.1
Inspecting a traced computation
For f1(x, y, z) = jnp.sum(x + y * z), jax.make_jaxpr(f1)(x, y, z) shows broadcasts, multiplication, addition, and reduction as separate equations. The result records shapes such as f32[3,3] and returns a scalar f32[]. This turns familiar array code into the sequence JAX transforms and compiles. section 5.2.1
Jaxpr records traced computation, not ordinary Python behavior. Printing occurs during tracing and is absent from the representation. Python loops and conditionals work when their behavior depends only on known values or shapes; input-dependent control flow can raise tracer conversion errors. Mark arguments static or use jax.lax structured control-flow primitives instead. section 5.1.2
