Three modes, one machine

The whole Python surface of eager mode is 32 lines: a setter, a getter, and a context manager that saves and restores. On the C++ side it is one boolean on the graph executor. Three call sites in the tensor layer read that boolean, and when it is set they call ApplyEagerSync on the tensors they just created.

ApplyEagerSync is one statement. It syncs the given tensors with wait false and sync_ltc_data false, which is to say it runs the same barrier machinery, per op, without turning everything into device data and without blocking. Same lowering, same builder, same LRU computation cache.

The cache part has an observable consequence the test suite pins. Run logsumexp on a 5 by 5 tensor twice and the eager compile counter stays at one while the eager execute counter goes to two. Eager mode is compile-per-op-shape, not compile-per-op. Multi-output ops go further: they construct every output with execution delayed, then call ApplyEagerSync once, so a multi-output op costs one compilation rather than N.

void XLAGraphExecutor::ApplyEagerSync(std::vector<XLATensorPtr>& tensors) {
  SyncTensorsGraph(&tensors, {}, /*wait=*/false, /*sync_ltc_data=*/false);
}

The repo publishes a measurement rather than an adjective, and since nothing in this course has a TPU attached, their measurement is the one to quote. A two-layer decoder-only model, fake data, 300 steps, one chip of a v4-8.

The caveat under the table matters more than the table. Pure eager reaches about 45 percent of compiled performance for decoder-only models, and about 1 percent for ResNet50. A number that swings by a factor of forty across two ordinary models is not a performance characteristic you can carry between workloads.

The docs draw the conclusion themselves: eager mode is for data preprocessing, random number generation, custom utilities and debugging, where immediate execution beats throughput. Not for the training loop, and not for the inference loop.

The decorator that recovers the lost throughput is a context manager, and its body is a sequence you can read in one screen. It turns eager off inside the region. It names the pending graph, syncs to flush whatever was already queued, then renames the graph to yours, so dataloading does not leak into the step graph and show up as a shape change.

It sets an allow_execution flag to the negation of full_graph. It optionally opens a dynamic-shape-detector session with a cap on how many distinct graphs the region may produce. Then it yields, and on the way out it restores everything and syncs again.

The full_graph=True enforcement is the part worth pausing on, because there is no strict-mode code path. When allow_execution is false, any execution inside the region reaches the debug analyzer, which prints an unexpected-execution report and then raises. The diagnostic is the enforcement mechanism. That also means the report you get names the Python frame that broke your full graph, which is more than a strict mode would have told you.

  @contextlib.contextmanager
  def _compile():
    saved_eager_mode_status = torch_xla._XLAC._get_use_eager_mode()
    saved_allow_execution = torch_xla._XLAC._get_allow_execution()
    saved_current_graph_name = torch_xla._XLAC._get_current_graph_name()
    torch_xla._XLAC._set_use_eager_mode(False)
    if name is not None:
      torch_xla._XLAC._set_current_graph_name(name + '_clear_pending')
    # Clear pending operations
    _clear_pending_ops_before_compile()

    if name is not None:
      torch_xla._XLAC._set_current_graph_name(name)

    # if full_graph sets to true execution can not happen before the sync below
    torch_xla._XLAC._set_allow_execution(not full_graph)

    if max_different_graphs is not None:
      torch_xla._XLAC._dynamic_shape_detector_set_max_different_graphs(
          max_different_graphs)
      torch_xla._XLAC._dynamic_shape_detector_start_session(current_id)

    try:
      yield
    finally:
      torch_xla._XLAC._set_allow_execution(saved_allow_execution)

Go looking for where torch.compile(backend="openxla") gets registered and you will not find it in pytorch/xla. There is no register_backend call for that name anywhere in the tree. Registration lives in PyTorch itself, and the definition is three lines that decide the shape of everything after.

openxla is aot_autograd wrapping a forward compiler. Because it goes through AOTAutograd, forward and backward reach the torch_xla bridge as separate FX graphs. That single fact is the source of the training-side cost, and the docs state the consequence outright: three graphs per training step instead of one.

What dynamo buys in exchange is real. Per step it skips FX interpretation, ATen dispatch, IR construction, the post-order walk, hashing and cache lookup by graph content, because the bridge holds a precomputed hash and jumps straight to execution. What it still pays per step is a Python loop splitting tensors from symbolic constants, an input materialization check that blocks if any input carries pending IR, an input list rebuild, and a copy back into every in-place-updated argument.

def openxla_eval_boxed(model, fake_tensor_inputs):
    return xla_backend_helper(model, fake_tensor_inputs, boxed=True)


openxla = aot_autograd(
    fw_compiler=openxla_eval_boxed,
)
register_backend(name="openxla", compiler_fn=openxla)

The lazy path can always recover from a cache miss by lowering the graph again. The dynamo path cannot, and the source says so with a TODO next to the check that fails.

ExecuteComputationWithBarrier looks its computation up by hash alone. There is no graph left to re-lower, because the whole point was to stop retracing. If the entry has been evicted from the 2048-entry LRU, the run throws. A long job with more distinct shapes than cache slots does not degrade; it stops.

Two smaller cliffs sit next to it. Under SPMD, dynamo cannot see an input sharding change, so the bridge compares shardings for a few steps and then stops checking; a sharding change after that crashes in the execution thread and leaves the result as a placeholder. And a dynamic-shape program does not retrace under dynamo but still compiles a fresh HLO per new shape, so the compile bill arrives even when the trace bill does not.

The lazy path can always re-lower. The dynamo path has thrown away the graph, so eviction is an error, not a recompile.

  // TODO implement a fallback mechanism, or make sure those entries
  // never get kicked out
  XLA_CHECK(cachedComputation)
      << "Failed to get computation by hash " << torch::lazy::HashToString(hash)
      << ". Maybe the entry get "
         "kicked out of the LRU cache";

The repo's own recommendation as of this commit splits by workload. Wrap the whole step function in torch_xla.compile for training. Use torch.compile(backend="openxla") for inference, where it lowers tracing overhead and the three-graph problem does not arise because there is no backward. The docs add that the long-term aim is for torch.compile to be the single compilation API for both.

Eager mode sits outside that split, as a debugging and utility mode. It is still torch_xla.experimental.eager_mode at this commit, and the docs say it is likely to become the default in future releases. The last lesson in this arc explains why that future release is probably not torch_xla.

Whichever mode you pick, the machinery under it is the one from lesson one. The hash rules do not change, the cut points do not change, and the seam under all three is the same class, which is where this arc goes next.