Guarded capture

Decorate a function with @torch.compile and dynamo attaches itself to it. On a call, dynamo walks the function's Python bytecode and captures every tensor operation it finds into an FX graph, the same kind of flat op list chapter seven reads directly. Alongside that graph it installs a set of guards: conditions on the shapes, dtypes, types, and closed-over values that were true when it captured. A later call that satisfies every guard reuses the cached graph outright. A call that breaks even one guard, a new shape, a different dtype, triggers a fresh capture and adds another entry to the cache.

dynamo compiles your tensor math and replays the rest.

Set this next to the jax path's model and the disagreement is exact, not a matter of degree. jax's jit traces a function once for a given signature, and the resulting jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → is the only thing that ever runs again; whatever Python ran during that one trace is gone from every later call. dynamo captures per guard set instead, as many times as your inputs actually vary, and it does something jax's tracer structurally cannot: the Python around the compiled tensor math keeps running on every single call, side effects included, not only once at capture time.

That last claim is worth proving rather than taking on faith, and the proof is small enough to run yourself. dict(counters["stats"]) is dynamo's own bookkeeping, a counter that tracks how many calls it has captured and how many distinct graphs exist. Call a compiled function f twice with the same shape and calls_captured reads 2 while unique_graphs stays at 1: two calls went through, one graph served both. Call it a third time with a new shape and both numbers move, to 4 and 2, because the guard on the first shape no longer matches and a second graph gets captured.

Notice what that sequence shows. calls_captured climbing on the second call, the one that hit the cache, means real work happened on that call beyond a lookup of a cached executable: enough of the wrapped function ran for dynamo's own instrumentation to fire again. Under jax's jit, a Python-level side effect like this would fire exactly once, during the single trace, and every call after that would be invisible to it, running only the recorded jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → with none of the original Python left around it. Under torch.compile, the side effect survives every call, cached graph or not, because the surrounding Python was never thrown away in the first place.

import torch
from torch._dynamo.utils import counters

@torch.compile(backend="eager")
def f(x):
    return torch.sin(x) + 1

f(torch.randn(4)); f(torch.randn(4))
print(dict(counters["stats"]))   # {'calls_captured': 2, 'unique_graphs': 1}
f(torch.randn(5))
print(dict(counters["stats"]))   # {'calls_captured': 4, 'unique_graphs': 2}

A jitted jax function facing a data-dependent branch raises outright, TracerBoolConversionError, because a tracer with no concrete value has nothing an if statement can test. dynamo has no sibling to that error. Write if x.sum() > 0: inside a compiled function with x traced, and dynamo does not fail. It stops capturing at that point, runs the branch as ordinary, uncompiled Python, and resumes capturing on the other side, stitching the compiled pieces around the gap.

That gap has a name: a graph break. Nothing about it raises or warns by default, which is exactly what makes it worth checking for on purpose rather than assuming it away. torch._dynamo.explain(g)(*args) runs the function and reports, among other things, a graph_break_count. Run it on the branch below and it comes back 1, one break, for the one condition that depends on a traced value.

The cost of a graph break is real but silent: instead of one large compiled graph, you get several smaller ones with an eager-mode gap between them, and every one of those boundaries is a place fusionSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → and codegen can't see across. A single break rarely matters. A compiled model with many of them starts looking, performance-wise, a lot like the eager code it was meant to replace, and explain is how you find that out before a slow benchmark tells you instead.

import torch
import torch._dynamo as dynamo

def g(x):
    if x.sum() > 0:        # data-dependent: a graph break, not an error
        return x + 1
    return x - 1

print(dynamo.explain(g)(torch.ones(3)).graph_break_count)   # 1

Everything in this chapter so far, the FX graph, the guards, the cache, the break, is dynamo's job: capture. What happens to a captured graph afterward, how it actually gets compiled into fast code, is a separate job, handled by a backend. Inductor is the default, and it's the one doing real codegen, choosing fusionsSeveral ops compiled into one kernel so intermediates stay in fast memory instead of round-tripping through HBM. XLA’s central optimization, with an exact limit.taught in /l/xla → and generating device code from the graph dynamo handed it.

backend="eager" is the setting used in this chapter's own examples, and it exists to isolate capture from codegen on purpose: with it, dynamo still captures, still installs guards, still counts calls and graphs exactly as shown above, but the captured graph just runs through PyTorch's ordinary eager dispatch instead of being compiled further. It's the right tool for studying dynamo's behavior in isolation, which is what this chapter needed, and the wrong one for measuring speed, which chapter nine's harness is built for instead.