What is a lazy tensor in PyTorch/XLA?

A tensor on the xla device that holds no numbers. Instead of a buffer it holds a pointer into a graph of operations that have not run yet. The C++ object behind it, XLATensor, can be in one of four states: it has a handle to memory already on the device, it has an IR value (a node in a recording of pending operations), it has a plain host tensor waiting to be uploaded, or it is a view of another tensor. When you write y = torch.relu(x @ w + b) on the xla device, PyTorch dispatches each operation as usual, but the XLA kernel builds an IR node instead of launching anything, and hands you back a tensor whose entire content is "I am the output of this node". Nothing computes until something forces it to.

What does torch_xla.sync() actually do?

It cuts the recording into a program and hands the program to the compiler. Concretely: sync() calls into the C++ extension, which collects every live tensor on the device that still has pending IR (not just the one you care about, every one still reachable from Python), sorts them into parameters that already have device data and body nodes that do not, folds their IR hashes into one graph hash, looks that hash up in a compile cache, compiles on a miss, and launches. It is a barrier over the whole device, not a flush of one tensor. That is why a tensor you kept in a list for logging ends up inside the same compiled program as your training step.

Why does PyTorch/XLA recompile when my input shape changes?

Because shape is part of the graph hash, and it gets there through the leaves. An ordinary IR node hashes its operation kind and folds in its operands, with no shape involved, so an eight-row matmul and a sixteen-row matmul hash the same. A leaf is different: a tensor that already has data on the device is represented by a DeviceData node, which uses a constructor that hashes the operation kind together with the shape string. Change the batch size and every leaf hash changes; because each interior node folds its operands in, that change propagates all the way to the root. New root hash, new graph hash, cache miss, full compile. This is why ragged batches or variable sequence lengths can compile hundreds of times, and why padding or bucketing to a small fixed set of shapes is the standard fix.

Why does adding a print() make my torch_xla training loop slow?

Because printing needs a real number, and getting a real number forces a sync in the middle of your step. print(loss), .item(), .cpu(), and float(t) all route to the same place, which compiles and runs whatever is pending for that tensor and then blocks on a device-to-host transfer. You pay a compile, a launch and a round trip inside the step. There is a second, subtler cost: that read path hashes the graph with an internal flag (force_ltc_data) set to false, while the sync-point path leaves it true, and that flag is the first thing folded into the graph hash. So the same operations reached through the two paths produce two different keys and two separate cache entries. A one-off debug print can leave you compiling the same logical graph twice. The fix is to keep device reads out of the hot path: read the loss every N steps, or accumulate on the device and read once per epoch.

How do I find out where my torch_xla graph is getting cut?

Two instruments, in order. First, torch_xla.debug.metrics.metrics_report() gives you CompileTime, ExecuteTime, and the UncachedCompile and CachedCompile counters. If the number of compiles is larger than the number of distinct input shapes you expected, something is cutting your step. Second, wrap the step in torch_xla.compile(full_graph=True). That sets an internal allow_execution flag to false for the region, and if any execution happens inside it, torch_xla prints the Python frames responsible and raises "Unexpected execution happens inside the compiled function, exiting". That turns a mysterious slowdown into a stack trace pointing at the offending line. There is a companion guard for the other failure mode too: max_different_graphs errors if a region ever produces more distinct traced graphs than the budget you set.

The Lazy Tensor: What Happens Between Your Op and sync()

On the xla device an op runs nothing. It appends a node to a graph, and the tensor you get back is a promise. This walks the whole mechanism in PyTorch/XLA at a pinned commit: what an XLATensor actually holds, how each IR node hashes itself as it is built and why shape only enters through the leaves, what torch_xla.sync() sweeps up and folds into one graph hash, the four ingredients of that hash in the order the executor folds them, the 2048-entry compile cache a hit or a miss lands in, and the reads (.item(), .cpu(), a stray print) that cut the graph somewhere you did not ask. Then the same machinery seen twice more: eager mode as a cut after every op, and torch.compile with the openxla backend as a cut pinned to the function boundary, traced once and replayed by hash.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of The Lazy Tensor: What Happens Between Your Op and sync() — built to be understood, not skimmed.

Questions

What is a lazy tensor in PyTorch/XLA?: A tensor on the xla device that holds no numbers. Instead of a buffer it holds a pointer into a graph of operations that have not run yet. The C++ object behind it, XLATensor, can be in one of four states: it has a handle to memory already on the device, it has an IR value (a node in a recording of pending operations), it has a plain host tensor waiting to be uploaded, or it is a view of another tensor. When you write y = torch.relu(x @ w + b) on the xla device, PyTorch dispatches each operation as usual, but the XLA kernel builds an IR node instead of launching anything, and hands you back a tensor whose entire content is "I am the output of this node". Nothing computes until something forces it to.
What does torch_xla.sync() actually do?: It cuts the recording into a program and hands the program to the compiler. Concretely: sync() calls into the C++ extension, which collects every live tensor on the device that still has pending IR (not just the one you care about, every one still reachable from Python), sorts them into parameters that already have device data and body nodes that do not, folds their IR hashes into one graph hash, looks that hash up in a compile cache, compiles on a miss, and launches. It is a barrier over the whole device, not a flush of one tensor. That is why a tensor you kept in a list for logging ends up inside the same compiled program as your training step.
Why does PyTorch/XLA recompile when my input shape changes?: Because shape is part of the graph hash, and it gets there through the leaves. An ordinary IR node hashes its operation kind and folds in its operands, with no shape involved, so an eight-row matmul and a sixteen-row matmul hash the same. A leaf is different: a tensor that already has data on the device is represented by a DeviceData node, which uses a constructor that hashes the operation kind together with the shape string. Change the batch size and every leaf hash changes; because each interior node folds its operands in, that change propagates all the way to the root. New root hash, new graph hash, cache miss, full compile. This is why ragged batches or variable sequence lengths can compile hundreds of times, and why padding or bucketing to a small fixed set of shapes is the standard fix.
Why does adding a print() make my torch_xla training loop slow?: Because printing needs a real number, and getting a real number forces a sync in the middle of your step. print(loss), .item(), .cpu(), and float(t) all route to the same place, which compiles and runs whatever is pending for that tensor and then blocks on a device-to-host transfer. You pay a compile, a launch and a round trip inside the step. There is a second, subtler cost: that read path hashes the graph with an internal flag (force_ltc_data) set to false, while the sync-point path leaves it true, and that flag is the first thing folded into the graph hash. So the same operations reached through the two paths produce two different keys and two separate cache entries. A one-off debug print can leave you compiling the same logical graph twice. The fix is to keep device reads out of the hot path: read the loss every N steps, or accumulate on the device and read once per epoch.
How do I find out where my torch_xla graph is getting cut?: Two instruments, in order. First, torch_xla.debug.metrics.metrics_report() gives you CompileTime, ExecuteTime, and the UncachedCompile and CachedCompile counters. If the number of compiles is larger than the number of distinct input shapes you expected, something is cutting your step. Second, wrap the step in torch_xla.compile(full_graph=True). That sets an internal allow_execution flag to false for the region, and if any execution happens inside it, torch_xla prints the Python frames responsible and raises "Unexpected execution happens inside the compiled function, exiting". That turns a mysterious slowdown into a stack trace pointing at the offending line. There is a companion guard for the other failure mode too: max_different_graphs errors if a region ever produces more distinct traced graphs than the budget you set.

The Lazy Tensor: What Happens Between Your Op and sync()

Questions

Related explainers