Does torch.compile with the openxla backend skip the lazy tensor?

No. There is no FX-to-HLO lowering anywhere in dynamo_bridge.py. The bridge calls the FX graph module on the caller’s real XLA tensors, and those operations dispatch to the same lazy backend they always do: they record an intermediate representation rather than execute. The bridge then asks the executor for the hash of that recording, tells it to compile under the hash without launching anything, restores every input the run modified from a clone taken beforehand, and clears the pending recording off every live tensor. What survives the trace is a hash, a parameter matcher, and two small output handlers. Every later call assembles the parameter list and calls _run_cached_graph with the saved hash, which looks the compiled program up in the compilation cache and executes it. So the lazy tensor is used exactly once per distinct set of inputs, and then never again for that graph.

Why does my compiled model keep recompiling?

Because something in your call pattern is producing a cache key the bridge has not seen. There are two caches. The compilation cache lives in C++ and is keyed by the graph hash, which folds in shapes, the operation sequence, and the build’s git revisions. The second cache is a plain Python dictionary held in the closure the bridge returns, and it is keyed by the tuple of non-tensor arguments in the call. When you compile with dynamic=True, or mark a dimension dynamic, Dynamo stops recompiling on a shape change and instead passes the varying dimension through as an ordinary integer argument. A new integer is a new key, a miss in that dictionary runs the whole tracing function again, and each trace ends in exactly one warm-up compile. A loop over a hundred distinct sequence lengths produces a hundred compiles, and the compilation cache holds 2048 entries by default before it evicts.

What happens to an operation XLA cannot lower inside a compiled region?

Nothing raises, and the region silently becomes several XLA programs. There is no static list of supported operations, so the bridge finds out empirically: it interprets the FX graph one node at a time, clearing the metric counters before each node and checking afterwards whether any operator fallback counter moved. A node that fell back, or whose arguments or results contain a tensor that is not on the xla device, goes on an unsupported list. That list becomes the predicate for PyTorch’s CapabilityBasedPartitioner, which groups the runs of supported nodes into fused submodules; each fused submodule is traced and compiled separately and spliced back into the graph as a call to a compiled closure. The unsupported nodes stay where they are and run through the ordinary dispatcher, which copies their tensors to the host, runs the CPU kernel, and copies the results back. Setting XLA_DYNAMO_DEBUG=1 prints the offending nodes.

Does the tracing pass actually execute my model on the device?

It runs your model, but it does not execute anything on the device. Calling the FX graph module records operations into the lazy runtime rather than computing them, and the compile step uses a warm-up-only sync that compiles the program into the cache with the wait flag off and no launch. What the tracing pass does do is mutate your tensors: an in-place operation in your model really rewrites the recording attached to the argument you passed. That is why the bridge clones every input tensor before tracing and copies the clones back afterwards, and why it clears the pending recording off every live tensor before returning. Skip either step and the next sync anywhere in your program would execute the traced graph a second time and apply the in-place updates again.

Why does the bridge clone all of my input tensors?

To undo the trace. Tracing here means really calling your model, so an argument the model writes to in place is genuinely changed by the time the trace finishes. The bridge takes a clone of every input tensor before the call and copies the clone back over the original afterwards, for exactly those arguments the runtime reports as needing materialisation. The restore has to happen before the pending recording is cleared, because the copy itself records operations and the clear is what removes them. The comment above the clone loop admits the clones cannot simply be traced instead, which would be the obvious fix: a clone shares device data with its base, and the device data records only the tensor id of the first tensor that pointed at it, so the mapping from data back to tensor id would name the wrong tensor.

torch.compile Meets the Lazy Tensor: dynamo_bridge.py, Line by Line

Name openxla as your torch.compile backend and PyTorch hands the captured FX graph to one 794-line Python file. The obvious guess is that it lowers those nodes to XLA directly. It does not. It runs the graph once, on your real device tensors, through the same lazy runtime an un-compiled program uses, takes the hash of the recording, compiles under that hash without executing, restores every tensor the run modified, deletes the recording, and hands back a closure that holds the hash. We read all of it: the matcher that rebuilds the parameter list from your arguments and the trace-time weights, the three small classes that put duplicated, pass-through and None outputs back before the caller sees them, the tracing function and its five undo steps, the closure that runs on every call and never compiles, the collector that discovers unsupported operations by running the graph node by node and watching a counter, and the partitioner that quietly turns one compiled region into several XLA programs with host code between them. The signature exhibit runs the file’s own branch: press a call, watch the numeric cache key decide whether the trace runs, and watch the launch happen either way.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of torch.compile Meets the Lazy Tensor: dynamo_bridge.py, Line by Line — built to be understood, not skimmed.

Questions

Does torch.compile with the openxla backend skip the lazy tensor?: No. There is no FX-to-HLO lowering anywhere in dynamo_bridge.py. The bridge calls the FX graph module on the caller’s real XLA tensors, and those operations dispatch to the same lazy backend they always do: they record an intermediate representation rather than execute. The bridge then asks the executor for the hash of that recording, tells it to compile under the hash without launching anything, restores every input the run modified from a clone taken beforehand, and clears the pending recording off every live tensor. What survives the trace is a hash, a parameter matcher, and two small output handlers. Every later call assembles the parameter list and calls _run_cached_graph with the saved hash, which looks the compiled program up in the compilation cache and executes it. So the lazy tensor is used exactly once per distinct set of inputs, and then never again for that graph.
Why does my compiled model keep recompiling?: Because something in your call pattern is producing a cache key the bridge has not seen. There are two caches. The compilation cache lives in C++ and is keyed by the graph hash, which folds in shapes, the operation sequence, and the build’s git revisions. The second cache is a plain Python dictionary held in the closure the bridge returns, and it is keyed by the tuple of non-tensor arguments in the call. When you compile with dynamic=True, or mark a dimension dynamic, Dynamo stops recompiling on a shape change and instead passes the varying dimension through as an ordinary integer argument. A new integer is a new key, a miss in that dictionary runs the whole tracing function again, and each trace ends in exactly one warm-up compile. A loop over a hundred distinct sequence lengths produces a hundred compiles, and the compilation cache holds 2048 entries by default before it evicts.
What happens to an operation XLA cannot lower inside a compiled region?: Nothing raises, and the region silently becomes several XLA programs. There is no static list of supported operations, so the bridge finds out empirically: it interprets the FX graph one node at a time, clearing the metric counters before each node and checking afterwards whether any operator fallback counter moved. A node that fell back, or whose arguments or results contain a tensor that is not on the xla device, goes on an unsupported list. That list becomes the predicate for PyTorch’s CapabilityBasedPartitioner, which groups the runs of supported nodes into fused submodules; each fused submodule is traced and compiled separately and spliced back into the graph as a call to a compiled closure. The unsupported nodes stay where they are and run through the ordinary dispatcher, which copies their tensors to the host, runs the CPU kernel, and copies the results back. Setting XLA_DYNAMO_DEBUG=1 prints the offending nodes.
Does the tracing pass actually execute my model on the device?: It runs your model, but it does not execute anything on the device. Calling the FX graph module records operations into the lazy runtime rather than computing them, and the compile step uses a warm-up-only sync that compiles the program into the cache with the wait flag off and no launch. What the tracing pass does do is mutate your tensors: an in-place operation in your model really rewrites the recording attached to the argument you passed. That is why the bridge clones every input tensor before tracing and copies the clones back afterwards, and why it clears the pending recording off every live tensor before returning. Skip either step and the next sync anywhere in your program would execute the traced graph a second time and apply the in-place updates again.
Why does the bridge clone all of my input tensors?: To undo the trace. Tracing here means really calling your model, so an argument the model writes to in place is genuinely changed by the time the trace finishes. The bridge takes a clone of every input tensor before the call and copies the clone back over the original afterwards, for exactly those arguments the runtime reports as needing materialisation. The restore has to happen before the pending recording is cleared, because the copy itself records operations and the clear is what removes them. The comment above the clone loop admits the clones cannot simply be traced instead, which would be the obvious fix: a clone shares device data with its base, and the device data records only the tensor id of the first tensor that pointed at it, so the mapping from data back to tensor id would name the wrong tensor.

torch.compile Meets the Lazy Tensor: dynamo_bridge.py, Line by Line

Questions

Related explainers