The anatomy of a jaxpr

Trace a JAX function and what comes back is a ClosedJaxpr, not a bare jaxpr. A jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →'s equations reference variables, and some of those variables are not arguments at all. They are values the traced function closed over, arrays or constants pulled in from outside the function body. A ClosedJaxpr bundles the jaxpr together with those captured values, called consts, so the whole thing can be evaluated, printed, or handed to a transform without any live Python closure behind it. When you print a traced function in a REPL, the ClosedJaxpr is what prints.

The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → itself has a small, fixed grammar. invars are the function's arguments, the variables a caller binds when it applies the jaxpr. constvars name the slots for closed-over values, matched positionally against the consts list on the ClosedJaxpr. Then come the equations, one per primitive the trace recorded, and finally outvars, the variables the jaxpr returns. Every equation binds a fresh variable on its left side and never rebinds it. That property has a name, and the last lesson in this unit spends itself on what the name buys you.

Each equation also carries params in square brackets, the primitive's static configuration, the part that is not itself a traced value. A dot_general equation's params hold its dimension numbers. A reduction's params hold the axes it reduces over. These are not extra arguments smuggled in; they are compile-time settings the primitive needs in order to run, fixed at trace time and baked into the equation rather than threaded through as a variable. Reading a jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → means reading both parts, the variables flowing through and the params sitting beside each equation.

The distinction between invars and constvars also tells you something about the function itself. Values in invars are whatever you called the function with, different on every trace. Values behind constvars are the ones fixed by the Python closure at trace time, most often weights or shapes baked in before tracing started. Two ClosedJaxprs with the same equations can carry completely different consts, which is why the consts list travels with the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → rather than living separately from it.

{ lambda ; a:bf16[32,64] b:bf16[64,64] c:bf16[64]. let
    d:bf16[32,64] = dot_general[
      dimension_numbers=(([1], [0]), ([], []))
      preferred_element_type=bfloat16
    ] a b
    e:bf16[1,64] = broadcast_in_dim[
      broadcast_dimensions=(1,)
      shape=(1, 64)
      sharding=None
    ] c
    f:bf16[32,64] = add d e
    g:bf16[32,64] = tanh f
  in (g,) }