One name, one definition

Write a function that assigns to the same Python name three times, x = x * 2.0, then x = x + 1.0, then x = jnp.tanh(x), and trace it on one f32[3] argument. Python had a single x the whole way down, rebound at each line. The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → has four variables.

Every equation binds a new letter on its left and no equation ever writes to a letter twice. mul binds b, add binds c, tanh binds d, and a, the argument, is never assigned again. Ask what x held at the second line and the printed program cannot answer, because there is no x in it. There is a value called b, and b means that value for the rest of the body.

The property has a name, and it is worth learning here because from this point down the stack you never meet an IR without it: static single assignment, SSA in every compiler paper you will read. Two rules, and the second one carries as much weight as the first. Each variable is bound by exactly one equation. Every use of a variable sits below the equation that bound it.

The word static is doing real work in that name. The rule is about the program text, not about how often a value gets computed while the program runs; a loop body binds its result once on the page and once per iteration on the machine. A jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → gets the second rule almost for free, since the body is a flat list in the order the trace recorded it, so reading top to bottom is reading in dependency order.

{ lambda ; a:f32[3]. let
    b:f32[3] = mul a 2.0
    c:f32[3] = add b 1.0
    d:f32[3] = tanh c
  in (d,) }

A property this load-bearing is worth checking rather than believing, and checking it takes about ten lines. Walk the equations, count left-hand sides by object identity so two different variables that happen to print as the same letter stay separate, and recurse into any jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → a primitive carries in its params so nested bodies are counted too.

Run it on the chapter gallery function dense and you get 4 bound, 4 distinct, 0 rebound. Run it on jax.grad of that function, a program with three times the equations, and you get 12 bound, 12 distinct, 0 rebound. A scan, whose body is a separate nested jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, comes back 4 bound, 4 distinct, 0 rebound. The count never diverges, on any program, which is what a structural guarantee looks like from the outside.

The other half of the property answers to a second pass. Walk the same grad program, keep a set of everything bound so far, and check each equation only reads names already in it: 12 equations, 0 uses before their definition. Nothing forward-references, so there is no order to reconstruct.

The letters themselves belong to the printer, not to the program. Two jaxprsThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → both using b are saying nothing to each other, which is why the count works on identity rather than on names.

from collections import Counter
from jax.extend.core import Jaxpr

def binders(jaxpr, counts):
    for eqn in jaxpr.eqns:
        for v in eqn.outvars:
            counts[id(v)] += 1              # one entry per left-hand side
        for p in eqn.params.values():
            for sub in (p if isinstance(p, (tuple, list)) else [p]):
                inner = getattr(sub, "jaxpr", sub)
                if isinstance(inner, Jaxpr):
                    binders(inner, counts)  # scan and cond bodies count too
    return counts

Write the same accumulation in C, a carry starting at zero and a loop adding one element at a time, and hand it to clang at -O1. The loop block that comes back opens with two instructions that have no counterpart anywhere in a jaxpr. %9 is the induction variable and %10 is the carry, and both are phi nodes.

The LLVM reference states what one does in a single sentence: at runtime, the phi instruction logically takes on the value specified by the pair corresponding to the predecessor basic block that executed just prior to the current block. Read the operands and the sentence explains itself. %10 is 0.000000e+00 if control arrived from block %4, the preheader, and %13 if it arrived from block %8, which is the block itself looping around.

The reason LLVM needs the instruction is structural. A function body there is a graph of basic blocks joined by edges, control can reach a block from more than one predecessor, and SSA still insists on one definition per name. So the merge itself has to be an instruction, and that instruction has to name the incoming edges. Nothing else in the language can express a value whose definition depends on how you got here.

A phi node is what SSA costs you when control flow is a graph of blocks.

A jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → body has no blocks and no edges to merge. It is one straight list of equations, and everything that would branch or repeat is a single equation holding whole jaxprs in its params. The next two sections read what that buys.

8:                                                ; preds = %4, %8
  %9 = phi i64 [ 0, %4 ], [ %14, %8 ]
  %10 = phi float [ 0.000000e+00, %4 ], [ %13, %8 ]
  %11 = getelementptr inbounds nuw float, ptr %0, i64 %9
  %12 = load float, ptr %11, align 4, !tbaa !6
  %13 = fadd float %10, %12
  %14 = add nuw nsw i64 %9, 1
  %15 = icmp eq i64 %14, %5
  br i1 %15, label %6, label %8, !llvm.loop !10

define float @accumulate(ptr noundef readonly captures(none) %0, i32 noundef %1) local_unnamed_addr #0 {
  %3 = icmp sgt i32 %1, 0
  br i1 %3, label %4, label %6

4:                                                ; preds = %2
  %5 = zext nneg i32 %1 to i64
  br label %8

6:                                                ; preds = %8, %2
  %7 = phi float [ 0.000000e+00, %2 ], [ %13, %8 ]
  ret float %7

8:                                                ; preds = %4, %8
  %9 = phi i64 [ 0, %4 ], [ %14, %8 ]
  %10 = phi float [ 0.000000e+00, %4 ], [ %13, %8 ]
  %11 = getelementptr inbounds nuw float, ptr %0, i64 %9
  %12 = load float, ptr %11, align 4, !tbaa !6
  %13 = fadd float %10, %12
  %14 = add nuw nsw i64 %9, 1
  %15 = icmp eq i64 %14, %5
  br i1 %15, label %6, label %8, !llvm.loop !10
}

Take a lax.cond on a scalar predicate, doubling on the true side and adding one on the false side. Three equations come back. gt binds c, the boolean, and convert_element_type binds d, turning that boolean into an int32 index. Then one cond equation binds e, and e is the answer whichever way the predicate goes.

Both branches sit in the params as nested jaxprsThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, each with its own binders, each ending in its own outvars. Neither one writes to e. They return, and the equation that called them binds the result. The merge that LLVM needed a phi instruction for is the ordinary act of a function returning a value into a caller that names it once.

Two details in this capture are worth pointing at. The branches are a tuple indexed by that int32, and index 0 is the false branch, which is why the add f 1.0 body prints first even though cond takes the true function first. And the printer draws every name from one running alphabet, so no letter repeats anywhere on the page, inside branch bodies included. The nesting lesson at /l/jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →/nested-jaxprs reads the indentation itself; the point here is only that the inner binders are separate from the outer ones and that the join has a single name.

{ lambda ; a:f32[] b:f32[3]. let
    c:bool[] = gt a 0.0
    d:i32[] = convert_element_type[new_dtype=int32 weak_type=False] c
    e:f32[3] = cond[
      branches=(
        { lambda ; f:f32[3]. let g:f32[3] = add f 1.0 in (g,) }
        { lambda ; h:f32[3]. let i:f32[3] = mul h 2.0 in (i,) }
      )
    ] d b
  in (e,) }

Now the accumulation the C loop was doing, written as a lax.scan: a scalar carry, an array of four elements, and a body that adds the element to the carry and returns it twice, once as the new carry and once as the per-step output. The whole loop is one equation.

The body is { lambda ; e:f32[] f:f32[]. let g:f32[] = add e f in (g, g) }, and num_carry=1 in the params says how many of its arguments are carried state. So the loop-carried value is e, a formal parameter of the body, bound fresh every time the body is applied. That is the same value the C version needed a phi for: zero on the first iteration, the previous %13 on every later one. One IR names the incoming edges, the other names an argument. The params lesson at /l/jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →/params-and-drills reads the rest of that bracket list.

The equivalence is old and it was stated plainly. Appel wrote a four-page argument in 1998 that SSA and functional programming are the same thing in different notation, and the sentence that matters for reading a scan is this one: the left-hand side of the phi assignment is the formal parameter of the corresponding function, and each right-hand side argument of the phi assignment is the actual parameter of some call to the corresponding function.

Wherever there is a formal parameter of a function (in the functional form), there is a phi (in the SSA form).

MLIR made the same choice one level below you, and says so in its rationale: regions represent SSA using block arguments rather than the phi instructions used in LLVM, a choice it calls representationally identical. The StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → lesson at /l/stablehlo/control-flow-is-regions reads a while carrying its state exactly that way, so the shape you are learning here survives the next lowering intact.

{ lambda ; a:f32[] b:f32[4]. let
    c:f32[] d:f32[4] = scan[
      _split_transpose=False
      jaxpr={ lambda ; e:f32[] f:f32[]. let g:f32[] = add e f in (g, g) }
      length=4
      linear=(False, False)
      num_carry=1
      num_consts=0
      reverse=False
      unroll=1
    ] a b
  in (c, d) }

Reading a large jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →, the question you ask most often is where a value came from. Under SSA that question is a search, not an analysis: the name appears on exactly one left-hand side, somewhere above the line you are staring at, and that equation is the whole answer. In a representation that allows rebinding you would have to reason about which assignment reached this point, and along which path.

The compiler gets the same discount, and the cheapest pass shows it best. Trace a function whose first two lines feed nothing, and the trace already prints the dead product as _, a variable nothing reads. Dead code elimination then walks the equations backward once, marking what is used, and drops the rest: the mul goes because nothing reads it, and the sin goes because the only thing that read it is gone. One reverse pass, no fixed point, because a use can only point upward.

The same reasoning is why the transforms in /l/jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →/transforms-rewrite can append and re-shape equations without ever renaming anything around them, and why every IR further down keeps the property rather than dropping it. StableHLOThe portable, versioned tensor IR that JAX and PyTorch both lower into; chapter 03 reads it line by line.taught in /l/stablehlo → and the MLIR it is written in keep it with block arguments, 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 → keeps it in HLO, where each instruction defines one value that later instructions reference by name. Learning to read equations this way once is learning to read the rest of the descent.

{ lambda ; a:f32[3]. let
    b:f32[3] = sin a
    _:f32[3] = mul b 3.0
    c:f32[3] = add a 1.0
    d:f32[3] = tanh c
  in (d,) }

after dce_jaxpr(jaxpr, [True]):

{ lambda ; a:f32[3]. let b:f32[3] = add a 1.0; c:f32[3] = tanh b in (c,) }