Part II — Front end

Capturing the graph

Lesson 01 priced eager execution and found two ceilings — launch-bound (thousands of tiny kernels, each ~3–10 µs to dispatch) and bandwidth-bound (each op re-reads its activation from HBM). Both are whole-program problems: to fuse two ops you must see both at once, to plan memory you must see every buffer's lifetime. But eager PyTorch hands the runtime one op at a time and forgets it the instant it returns. There is no program to optimize — only a stream of already-executed calls. Before any pass in this track can run, you have to do the unglamorous thing first: turn the Python into a graph. This lesson is about the three ways to capture one, and the real trade hiding in each.

1 · Why eager hands you no graph

In eager mode every tensor op is dispatched and executed the moment Python reaches it. z = x @ w walks Python → the PyTorch dispatcher → an ATen kernel → a CUDA launch → an output tensor, and then control returns to the interpreter. The next line repeats the whole walk. There is no object anywhere that says "z is the matmul of x and w, and it feeds the next add." That edge existed for a microsecond as a stack frame and is now gone. This is exactly what makes eager pleasant to debug — you can print any intermediate, set a breakpoint, branch on a tensor's value — and exactly what makes it impossible to optimize across ops.

A graph is the thing eager refuses to build: a directed acyclic graph whose nodes are ops and whose edges are tensor values, capturing which op produces which value and who consumes it. With that in hand a compiler can finally ask the lesson-01 questions — which round trips can I merge, which buffers can share storage, which layout cancels a transpose. Capture is the act of recovering that graph from the imperative program. There are exactly three places to intercept the program to do it, and they line up with how much of the program's structure each one can see.

2 · Tracing: run once, record the ops

Tracing runs the function on example inputs and records every tensor op the execution actually performs, in order, wiring outputs to inputs to build the graph. It is the simplest strategy and the most portable — the result is just a list of ops, which any backend can consume. torch.fx.symbolic_trace does it with Proxy objects (fake tensors that record operations instead of computing them); torch.jit.trace does it by running real tensors and watching the dispatcher; JAX's make_jaxpr traces a function into a jaxpr — JAX's typed, functional trace IR. A trace, then, is the recorded op sequence of one execution. Here is the jaxpr for a tiny linear-plus-gelu, abridged:

>>> import jax, jax.numpy as jnp
>>> def f(x, w, b): return jax.nn.gelu(x @ w + b)
>>> print(jax.make_jaxpr(f)(jnp.ones((8, 4)), jnp.ones((4, 4)), jnp.ones(4)))
{ lambda ; a:f32[8,4] b:f32[4,4] c:f32[4]. let
    d:f32[8,4] = dot_general[dimension_numbers=(([1],[0]),([],[]))] a b
    e:f32[8,4] = add d c
    f:f32[8,4] = ... gelu(e) ...           # exp / tanh / mul chain, abridged
  in (f,) }

Notice the shapes are frozen into the type of every value: f32[8,4], not f32[batch,4]. That is the first half of tracing's trap — the trace is specialized to the example's shapes. Feed it a batch of 16 and the trace is, in general, invalid; you re-trace. The second, more dangerous half is control flow. The tracer only ever sees the ops that executed, so a data-dependent branch collapses to whichever side the example happened to take:

def f(x):
    if x.sum() > 0:          # a Python bool on a traced tensor — the tracer can't see both sides
        return x * 2
    return x - 1

# trace with x.sum() > 0  →  graph is just  {mul x, 2}.  The  (x - 1)  branch is GONE.
# run the trace on an x with negative sum  →  you still get  x * 2.  Silently wrong.

The if is not in the graph because it was never a tensor op — it was a Python control decision that the tracer resolved once and baked in. fx.symbolic_trace at least errors if you call bool() on a Proxy (you can't branch on a fake tensor); jit.trace warns but proceeds. Either way, tracing's verdict is blunt: it captures dataflow perfectly and control flow not at all. For a feed-forward block with fixed shapes that is fine. For anything with a real branch — early exit, padding-dependent masking, mixture-of-experts routing — it is a correctness bug waiting to ship.

3 · Scripting / AST: parse the source, keep the branches

If the problem with tracing is that control flow lives in the source and never reaches the runtime, the obvious fix is to capture from the source. Scripting parses the Python AST and compiles it — if, for, while and all — into a graph IR that contains real control-flow nodes. torch.jit.script is the canonical example: it walks the function's AST and emits TorchScript, an IR with genuine branch and loop ops, so both sides of every if are preserved regardless of which one any particular input takes.

@torch.jit.script
def f(x):
    if bool((x.sum() > 0).item()):   # captured as a real prim::If node — BOTH branches live in the graph
        return x * 2
    return x - 1
# graph(%x): %1 : bool = aten::gt(aten::sum(%x), 0)
#            %y : Tensor = prim::If(%1)
#              block0(): -> (aten::mul(%x, 2))     # then
#              block1(): -> (aten::sub(%x, 1))     # else

The cost is paid in coverage. The AST compiler can only translate the subset of Python it has rules for: static types (it infers or you annotate; defaults to Tensor), a fixed set of builtins, no arbitrary *args games, no calling into NumPy or third-party libraries, no dynamic class tricks. Hit something outside the subset and scripting fails to compile — loudly, which is at least honest, but it means real models need rewriting to fit the subset. This is why torch.jit.script never became the default capture path: the subset was a constant tax on model authors, and the field wanted something that captured ordinary, unmodified PyTorch.

4 · Dynamo: interpret the bytecode, guard, and break

TorchDynamo (the capture front end of torch.compile) is the third strategy, and it splits the difference. Instead of running the function (tracing) or parsing its source (scripting), it hooks CPython's frame evaluation — the layer below the source — and symbolically interprets the bytecode. As it steps through the instructions, it accumulates the tensor ops into an FX graph (PyTorch's flat IR: a list of (op, args, kwargs) nodes with tensor metadata; see gpu_kernels · 23 torch.compile). When it reaches an instruction it can model as a tensor op, it adds a node. When it reaches one it cannot — a data-dependent Python branch, a print, a call into NumPy, a list mutation — it does a graph break: it compiles the straight-line graph accumulated so far, hands control back to the real CPython interpreter to run the un-traceable bit eagerly, then starts a fresh graph after. The result is a chain of compiled graph islands separated by eager gaps.

Because Dynamo captures a straight-line graph (it followed one path through the bytecode), it must protect that graph's assumptions. It does this with guards: a set of cheap runtime checks — on input dtypes, shapes (or symbolic shape relations), tensor ids, Python constants, the values it branched on — recorded alongside the compiled graph. Before reusing a cached graph, the runtime checks its guards; if every guard passes, it replays the compiled kernels, if any guard fails, the graph is invalid for these inputs and Dynamo recompiles. So Dynamo handles the cases tracing botched in two different ways: a data-dependent branch becomes a graph break (run that branch eagerly, capture each side's straight-line continuation separately), and a changed shape or dtype trips a guard and triggers a recompile rather than silently running a stale graph. A compact guard set looks like this:

# torch._dynamo guards recorded for one compiled graph (illustrative, abridged):
TENSOR_MATCH:  L['x']  dtype=torch.float32  device=cuda:0  rank=2  requires_grad=False
SHAPE_ENV:     L['x'].size()[0] == 8        # batch specialized to 8  (recompiles if it changes)
               L['x'].size()[1] == 4
CONSTANT_MATCH L['scale'] == 2.0            # a Python float captured by value
#  miss any guard  →  cache miss  →  retrace + recompile this frame

And the FX graph Dynamo emits for the straight-line part of our running example — captured, then fed to the rest of the compiler — reads like the jaxpr did, but as PyTorch nodes:

opcode         name     target                      args              # FX graph dump (torch.fx)
-------------  -------  --------------------------  ----------------
placeholder    x        x                           ()
placeholder    w        w                           ()
placeholder    b        b                           ()
call_function  mm       aten.mm                     (x, w)
call_function  add      aten.add                    (mm, b)
call_function  gelu     aten.gelu                   (add,)
output         output   output                      (gelu,)

This is the verdict the field reached: capture ordinary, unmodified PyTorch (no scripting subset to satisfy), keep correctness in the face of control flow and shape change (graph break + guard, never a silent wrong answer), and accept that the captured region may be fragmented. The remaining art — and most of lesson 11's job — is keeping graph breaks rare (each one stops fusion and CUDA-graph capture at the boundary) and guards loose enough to avoid a recompile storm.

5 · Drive it: the capture simulator

The widget runs a tiny program — one shape and one data-dependent branch — through each capture strategy. Pick a strategy, then perturb the world: flip the branch the input takes, or change the input's shape. Watch the four KPIs. Tracing always captures a graph but loses the branch (and goes stale when the shape changes, with no signal); scripting keeps the branch but only because the program stays in-subset; Dynamo keeps correctness by breaking the graph at the branch and by firing a guard → recompile when the shape changes. The configuration that "breaks" is tracing with the branch flipped: it reports a captured graph and preserved-looking output while having silently dropped the path the new input needs.

Failure modes & checklist

Checkpoint

Where this points next

Whichever strategy you pick, the output is a graph — a jaxpr, a TorchScript graph, an FX graph. But "a graph" here is still just a pile of framework objects: ATen op nodes, Proxy edges, Python constants, in-place mutations that alias each other, types that may or may not carry a shape. There are no rules — no statement of what counts as a value, what a legal rewrite is, why moving an op past another preserves the math. Try to run constant folding or fusion on this raw capture and you will trip over aliasing and mutation immediately. To rewrite a graph safely and repeatedly you need a representation with a contract: single-assignment values, types that carry shape and dtype and device, a small well-defined op set, and levels you can lower between. That is the intermediate representation, lesson 03 — The intermediate representation.

Interview prompts

• Why can't you optimize an eager program without capturing it first? (§1 — eager dispatches and executes each op immediately and discards the dataflow relationship; cross-op transforms like fusion and memory planning need the whole producer/consumer DAG, which only exists once you capture a graph.)
• What does tracing silently get wrong, and why? (§2 — it records only the ops that executed, so a data-dependent if collapses to the path the example took and the other branch never enters the graph; shapes are also frozen into value types. No tensor op = nothing to capture.)
• How does scripting preserve control flow, and what does it cost? (§3 — it parses the AST and emits real branch/loop ops (prim::If), so both sides survive; the cost is that it only compiles a typed Python subset, so unsupported constructs fail to compile.)
• What exactly is a guard, and what happens when one fails? (§4 — a cheap runtime check on a captured graph's assumptions (dtype, shape/symbolic-shape relation, tensor id, captured constant); on failure the cached graph is invalid for these inputs, so Dynamo retraces and recompiles the frame.)
• What is a graph break and what does it cost you? (§4 — Dynamo hits an instruction it can't model, compiles the graph so far, runs the un-traceable part in eager CPython, then starts a new graph; it bounds the captured region, so fusion and CUDA-graph capture stop at the boundary.)
• You torch.compile a model and see a recompile every step. Diagnose. (§4 / Failure modes — a guard is over-specialized on a value or shape that changes each call; loosen it / mark the dim dynamic so the graph is reused instead of retraced — full treatment in lesson 11.)
• Which capture strategy would you pick for a feed-forward block vs a model with early-exit, and why? (§5 — tracing for the branch-free block (cleanest, most portable graph); Dynamo (or torch.cond) for early-exit, because a bare trace would drop the exit path and ship a correctness bug.)