greatitself.
Continue
advanced27 min studyLesson 511 in Engineering

Capture a tensor region without confusing it with the whole program

Separate tracing, guarded reuse, eager execution, and ahead-of-time export using a small inference pipeline and an explicit compilation budget.

Before you begin
On this page

The question for today

What you’ll understand

  • Identify the tensor operations and host behavior included in a captured region.
  • Distinguish partial JIT capture from a full export contract.
  • Calculate when a fictional compilation cost is amortized.

A tensor compiler does not begin with an arbitrary Python program already expressed as one optimizable array formula. It must establish which operations, values, and assumptions belong to a captured region. A serving function can multiply a batch, inspect a result, update a host counter, and then launch another tensor operation. Each boundary changes what the compiler can analyze together. The version-sensitive examples here use PyTorch 2.9 and JAX 0.8.0 documentation; they illustrate contracts rather than promising identical capture behavior in every later release.

Trace execution, then state what was captured

Imagine an inference function receiving X with shape [B,4] and W with shape [4,3]. The tensor region computes a matrix product and adds a three-element bias. Surrounding Python might validate a request identifier, format a log line, and choose a response envelope. Capturing the numerical region does not automatically capture those host operations. A graph should identify its inputs, outputs, operation semantics, and relevant assumptions, so a reviewer can tell whether a missing operation was intentionally external or accidentally omitted.

PyTorch Dynamo observes Python execution and extracts tensor graphs while arranging guarded reuse. In ordinary partial-capture compilation, unsupported portions can create graph breaks with eager execution between captured regions. Export has a different goal: obtain a complete representable tensor program under its declared constraints, and reject unsupported capture rather than silently preserve arbitrary Python around it. These are useful alternatives for different deployment needs. Neither a graph drawing nor a successful trace alone demonstrates that a particular backend can lower every operation efficiently.

Y=XW+1BbT,X∈RB×4, W∈R4×3, b∈R3Y=XW+\mathbf1_B b^{\mathsf T},\qquad X\in\mathbb R^{B\times4},\ W\in\mathbb R^{4\times3},\ b\in\mathbb R^3
The ideal real-valued region has output shape [B,3], with positive integer B and an explicit broadcast over rows. Actual tensor dtypes, rounding, device, and operator definitions are additional parts of the execution contract; shape compatibility alone is insufficient.↔ Scroll wide equations horizontally.

A trace is valid only within its assumptions

A branch based on a known shape can be specialized and protected by an input condition. A branch based on an actual tensor value is different: the example value observed while tracing is not generally a constant for future requests. Supported control-flow operators, explicitly static values, or an eager boundary may be appropriate, depending on the capture system. Do not delete the branch merely because the first batch took one path. The alternative path is a direct candidate for a counterexample test.

Guard failure means the existing artifact's assumptions do not hold for that call; it is not evidence that the input program is inherently wrong. A runtime may select another valid artifact, capture a new one, or execute eagerly according to its policy. Record capture events separately from backend compilation and runtime execution. Otherwise a benchmark can attribute repeated tracing time to matrix multiplication, or accidentally compare a warm compiled path with an eager path that includes unrelated initialization.

Tstaged=C+NR,Teager=NET_{\rm staged}=C+NR,\qquad T_{\rm eager}=NE
C is one compilation cost in milliseconds, R and E are per-call times in milliseconds, and N is a positive integer call count. This simplified budget assumes one reusable artifact, no additional recompilations, and comparable synchronized work. It is a fictional accounting model, not a measured speed claim.↔ Scroll wide equations horizontally.

Reuse is a workload property worth testing

An artifact can be profitable when many requests satisfy the same contract and unprofitable when every request triggers a fresh specialization. Therefore test both the expected production shapes and deliberate changes in metadata or static arguments. Keep a log of why reuse failed, rather than replacing all inputs with one convenient benchmark shape. Export constraints and runtime guards should be inspectable evidence connecting an optimized region to the program it is supposed to implement.

The cache lab offers a bounded analogy for artifact reuse, not an implementation of Dynamo's cache policy. Reset, set both capacity C=4 and request stride s=4, predict twelve hits among the sixteen requests already evaluated, then Play or Step. Reset and set both C=3,s=4 for a separate zero-hit comparison. Four repeating identifiers fit only in the first capacity. Playback extends this deterministic LRU stream to at most 4,096 requests; it performs no tracing, guard evaluation, or compiler invocation.

Work it through

Amortize one captured region

A fictional pipeline takes E=6 ms per eager call. Capturing and compiling one reusable region costs C=120 ms, after which a comparable call takes R=4 ms. Assume all calls satisfy its guards and timings include required completion waits.

  1. Compare cumulative budgets rather than just warm execution: eager costs 6N ms and the staged route costs 120+4N ms. Both cover the same completed work.

  2. Solve 120+4N<6N, giving N>60. At exactly sixty calls the budgets tie at 360 ms; the first whole-number call count with a modeled saving is sixty-one.

  3. For one hundred calls, eager costs 600 ms and the staged route 520 ms, saving 80 ms under these assumptions. A second compilation or an eager boundary would require an additional budget term.

The result

The modeled benefit begins at sixty-one compatible calls. This calculation does not establish that an actual backend achieves the assumed four-millisecond call time.

The living notebook

Explore the model

Interactive lab

Interactive lab · cache

Inspect repeated identities before assuming compilation reuse

Use fresh Reset → set both C=4,s=4 and Reset → set both C=3,s=4 comparisons. Predict twelve versus zero initial hits, then Play or Step. Sixteen requests exist at time zero; the stream caps at 4,096. This LRU model does not run a compiler or implement its guards.

Inspect repeated identities before assuming compilation reuseA fully associative LRU cache starts empty and requests address (i·s) mod 16. Hits move an address to the newest position; misses evict the oldest when full. Playback starts with 16 requests and adds two per second, capped at 4096 explicitly simulated requests. This toy stream is not a processor-cache benchmark. Current readings: Hits / requests: 0 / 16; Hit rate: 0%; Address-stream period: 16 addresses. All parameters are adjustable with the labeled controls below.Fully associative LRU cache after 16 deterministic requests: 0 hits and 16 misses.Most recent 16 requestsaddress = i × 3 mod 160M3M6M9M12M15M2M5M8M11M14M1M4M7M10M13MCurrent cache order · oldest → newest471013Latest address 13 · miss: inserted as newestToy address stream · period 16 · starts empty
Cached addressMiss / evictionMost recent request
Hits / requests
0 / 16
Hit rate
0%
Address-stream period
16 addresses
t = 0 s

A fully associative LRU cache starts empty and requests address (i·s) mod 16. Hits move an address to the newest position; misses evict the oldest when full. Playback starts with 16 requests and adds two per second, capped at 4096 explicitly simulated requests. This toy stream is not a processor-cache benchmark. Parameter changes pause playback and recompute the model at the current elapsed time. Reset restores both parameters and returns time to zero. Playback starts paused. Adjust either parameter, then choose Play or Step.

Look forRepeated valid artifact use is a workload assumption that needs evidence.

Pause & reason

Check your understanding

A future input takes a tensor-value branch absent from the initial trace. Which response preserves a sound capture contract?

Your turn

Try it yourself

The same pipeline needs another 120 ms compilation after fifty calls because its assumptions change. Across one hundred calls, suppose execution still costs 4 ms per call and eager costs 6 ms. Does the earlier eighty-millisecond saving survive?

Take a moment to make a prediction or work it out on paper. Open a hint when you need one.

Give me a hint

Count both compilation events before comparing the two cumulative totals. Reuse of one artifact was an explicit premise of the first calculation.

Compare with the solution

The staged total is now 2×120+100×4=640 ms, while eager remains 100×6=600 ms. The modeled staged path is forty milliseconds slower. This does not make graph capture universally undesirable; it identifies specialization churn as a missing cost in the original budget. Inspect the changed assumptions and workload before deciding whether a more general artifact or selective eager execution is preferable.

Keep the idea

Take it with you

  • Capture boundaries determine which behavior the compiler can analyze together.
  • An artifact is reusable only while its declared assumptions hold.
  • Compilation profitability depends on compatible call counts and all boundary costs.

Read further

Original sources and deeper explanations to follow at your own pace.

Optional progress, saved in this browser when available. Revisit any lesson at any time.

Lesson not marked complete.