Read your first jaxpr

Read a jaxpr as a typed computation

Start at the inputs and trace one intermediate to its consumers. For a reduction, identify the consumed axes and resulting shape. Distinguish array inputs from closed constants: both can affect the output, but they enter the program differently.

Then connect each primitive to the corresponding part of the original expression. A variable name in a printed jaxpr is a local label, not a stable API to test verbatim.

The intermediate-value heatmap is a numerical aid. Pair it with the typed dependency graph to explain how those values were produced. Lowering and compilation can change representation later; a primitive count alone does not predict executed kernel count or runtime.

Compute the values before reading the representation

Start with $x=[1, 2, 3]$. Adding $1$ gives $[2, 3, 4]$. Elementwise multiplication of this vector by itself gives $[4, 9, 16]$. The sum is $29$. The first two operations retain shape $(3,)$, while the reduction produces a scalar, whose array shape is ().

The derivative of each $(x_i+1)^2$ is $2(x_i+1)$, so the gradient at this input is $[4, 6, 8]$. A different implementation could produce the same values with different primitives. We use multiplication explicitly to keep this first trace readable. Do not infer an algorithm from the spelling of a Python expression alone.

x: f32[3] = [1, 2, 3]
   │ add scalar 1
   ▼
shifted: f32[3] = [2, 3, 4]
   │ multiply by shifted
   ▼
squared: f32[3] = [4, 9, 16]
   │ reduce_sum over axis 0
   ▼
result: f32[] = 29

Read one equation at a time

In a typical printout the input is named a, the shifted vector b, the squared vector c, and the scalar result $d$. These names are generated; they are not your Python variable names and can change across versions. An equation such as b:f32[3] = add a 1.0:f32[] binds a new value rather than mutating a. f32 means float32, $[3]$ denotes a length-three vector, and [] denotes a scalar.

The lambda header lists closed constants before its semicolon and ordinary input variables after it. The let block lists equations; the final in (...) lists returned values. Parameters in brackets belong to a primitive, for example reduction axes. Our script iterates over eqns and prints primitive names and output shapes so you can relate the structured object to the text. Treat exact formatting and generated names as version-dependent details.

{ lambda ; a:f32[3]. let
    b:f32[3] = add a 1.0:f32[]
    c:f32[3] = mul b b
    d:f32[] = reduce_sum[...] c
  in (d,) }
Illustrative excerpt: actual reduction parameters may differ.

Trace an abstract signature, not the sample’s values

make_jaxpr uses an example array to establish an abstract input shape and dtype. It does not use the ordinary data values as Python control-flow decisions. Calling the same function with another float32 length-three array follows the same computation structure. A length-four input changes the types in the representation, even though the algorithm is still add, multiply, sum.

This is why if $x$$[0]$ > $0$ cannot select a Python branch during ordinary tracing: the comparison produces a traced boolean, not a concrete Python truth value. Use array-compatible control flow such as lax.cond with compatible branch result types. Moving the predicate into a static argument is a different design and may trigger separate traces; it is not a universal repair.

Inspect the gradient as another program

make_jaxpr(grad(shifted_square_sum))(x) asks JAX to trace the transformed derivative function. You will see additional operations that propagate a scalar output sensitivity back to vector inputs, including multiplication and broadcasting. Some forward operations may remain; their exact arrangement depends on transformations and JAX’s implementation.

Count the equations as a descriptive observation, then verify the numerical derivative independently at multiple inputs. More equations do not automatically mean slower execution: lowering, fusion, memory movement and hardware scheduling affect runtime. A jaxpr can tell you what was traced; profiling answers where time is spent.

Separate data inputs from captured constants

An array defined outside a function and read by it can appear as a closed constant. In the exercise below the offset vector is $[1, 2, 3]$; $x$ remains the ordinary input. Inspect closed.consts and jaxpr.constvars, then compare with a version that takes the offset as an explicit second argument. The explicit version has two input variables.

This distinction matters for clarity and reusable function signatures. It does not prove that a constant occupies a particular device memory buffer or is copied on every call. Later compiler stages can represent constants differently. Save the traced structure without presenting it as a memory trace.

Compare Python unrolling with structured iteration

A Python loop with a fixed range is executed during tracing, so its repeated array operations can appear repeatedly in the outer jaxpr. lax.scan represents a loop with a nested body and a carry contract; the outer equation list includes the loop primitive rather than simply expanding every iteration.

Inspect both representations for a four-step recurrence and verify they return the same result. The body jaxpr still contains operations, so counting only outer equations would hide work. This comparison illustrates representation, not a benchmark. For an actual performance claim, measure compilation and synchronized execution separately on the intended hardware.