The seam

ComputationClient is a pure-virtual class in torch_xla/csrc/runtime/computation_client.h, and it exists because the LazyTensor core above it deals in opaque handles. Its nested Data inherits the core's backend data type; its nested Computation inherits the core's computation type. The class re-types those generic handles into XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →-shaped ones carrying an xla::Shape and an xla::OpSharding, and it hides whether the bytes underneath are a PJRT buffer or an IFRT array.

Two subclasses implement it. There is no registry and no dynamic dispatch beyond the vtable; one function picks the implementation once, at first use. Read that function and the choice turns out not to be a choice at this commit. The IFRT client is compiled into the binary, it is a dependency of the runtime build target, and it is unreachable, because the environment variable that used to select it is commented out and the boolean above the branch is a hard-coded false.

What that parked client would have changed is worth knowing without retelling: an IFRT array carries its own sharding across devices where PJRT hands out per-device buffers. The xla path's lesson at /xla/ifrt/above-the-compiler is where that comparison lives. What the torch_xla copy adds is a list of what does not work yet. It refuses non-SPMD compilation outright, ExecuteComputation returns unimplemented, and roughly a dozen methods throw, serialization among them, which means no persistent compilation cache.

  // TODO: enable IFRT once it's not crashing anymore.
  // Ref: https://github.com/pytorch/xla/pull/8267
  //
  // static bool use_ifrt = sys_util::GetEnvBool("XLA_USE_IFRT", false);
  const bool use_ifrt = false;
  if (sys_util::GetEnvString(env::kEnvPjRtDevice, "") == "") {
    return XLA_ERROR_WITH_LOCATION(
        absl::FailedPreconditionError("$PJRT_DEVICE is not set."));
  }

  ABSL_CHECK(!g_computation_client_initialized)
      << "ComputationClient can only be initialized once.";

  std::unique_ptr<ComputationClient> client;
  if (use_ifrt) {
    XLA_ASSIGN_OR_RETURN(client, IfrtComputationClient::Create());
  } else {
    XLA_ASSIGN_OR_RETURN(client, PjRtComputationClient::Create());
  }

Follow one tensor from a torch op to device memory and back and you cross this interface thirteen times, in thirteen distinct places. The table below pairs each frontend site with the client method it calls, at file and line, so you can check every row rather than take the count on faith.

Read the middle column and the shape of the interface shows itself. Compile and execute, transfers in both directions, a placeholder allocator for the asynchronous path, an environment hash, serialization for the persistent cache, two topology queries and one sharding query. That is a runtime contract, not a device API. Nothing in it mentions a stream, a kernel, or a collective.

One honest caveat about the count: the topology queries in particular are called from more places than this table lists, including the Python bindings and the ATen bridge. The table follows one tensor's path rather than inventorying every call site in the tree.

Two of those thirteen rows sit next to each other in one if-else, on a worker thread, inside the lambda the executor schedules. Whether SPMD is on decides which arm runs, and the arms differ in more than a name: the replicated arm hands the client a device list and gets back sharded data, and the single-device arm hands it one device string and options carrying the eager-mode flag from lesson three.

Neither arm blocks. Both return handles to buffers that may not be ready, which the executor then assigns onto placeholders created earlier. Python returns from sync() with nothing computed and nothing waited on. The first blocking call comes later, when something reads a value, and it lands on TransferFromDevice, whose header comment carries a warning worth internalizing: it blocks until the data is ready, and calling it from Python while holding the GIL can deadlock.

The order for a full step, then, reads: compile if the hash missed, transfer any host-side parameters up, schedule the execute on a worker thread, assign results into placeholders, return. Read a value and only then wait. Everything about this pipeline is designed so the Python thread is never the thing waiting.

        XLA_ASSIGN_OR_THROW(
            std::vector<runtime::ComputationClient::DataPtr> outputs,
            client->ExecuteReplicated(*async->cached_computation->computation,
                                      UnwrapXlaData(async->parameters_data),
                                      devices, execute_options));
        results = WrapXlaData(outputs);
        TORCH_LAZY_COUNTER("ExecuteReplicated", 1);
      } else {
        XLA_ASSIGN_OR_THROW(
            std::vector<runtime::ComputationClient::DataPtr> outputs,
            client->ExecuteComputation(*async->cached_computation->computation,
                                       UnwrapXlaData(async->parameters_data),
                                       async->device.toString(),
                                       {/*explode_tuple=*/true,
                                        /*eager_mode=*/use_eager_mode}));
        results = WrapXlaData(outputs);
        TORCH_LAZY_COUNTER("ExecuteComputation", 1);

Look for a donation list on the execute call and there is none. torch_xla never puts one there. Which parameters may be donated is decided before compilation and written into the module itself, by calling AddBufferDonor on the XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → builder for each donor index, so the compiled executable carries the input-output alias as part of its own definition.

PJRT's side of that arrangement is stated in the class comment above PjRtLoadedExecutable, and the two quotes below are worth reading as a pair. One side writes the alias into the computation; the other side reads it off the computation and donates the parameter at execution time. The runtime is never told; it is compiled at.

There is a consequence for the hash from lesson one. Because donation changes the module, the donor indices are merged into the graph hash, and they are merged as a sorted set specifically so the ordering cannot make the hash unstable. There is a matching consequence for parameter tupling: when a program crosses 3200 parameters, all parameters get wrapped into one tuple, and the donors are re-emitted against that tuple by shape index instead of by parameter number.

What donation buys you, and what it costs the donated array, is the xla path's story and it is told at /xla/layout-memory. The opt-out on the PJRT side, the non-donatable input indices field on ExecuteOptions, is told in the xla path's chapter 14. The torch_xla-specific fact is only this: that field is never the mechanism here.

// torch_xla/csrc/xla_graph_executor.cpp: the alias goes into the module
void XLAGraphExecutor::SetBufferDonors(
    LoweringContext* lowering_ctx,
    const std::vector<size_t>& buffer_donor_indexs) {
  for (size_t i : buffer_donor_indexs) {
    lowering_ctx->builder()->AddBufferDonor(/*param_number=*/i,
                                            /*param_index=*/{});
  }
  TORCH_LAZY_VALUE_METRIC("InputOutputAliasCount", buffer_donor_indexs.size());
}

// xla/pjrt/pjrt_client.h: and PJRT reads it back off the module
// Represents a compiled computation that can be executed given handles to
// device-allocated literals. If any input/output alias has been specified in
// the computation, the parameter containing the input buffer will be donated
// when passed to the execution.
class PjRtLoadedExecutable {

Two absences from the interface say as much as the thirteen crossings. There is no collective call on it at all. all_reduce and its relatives lower into the graph as HLO collectives, so by the time anything reaches the runtime the communication is already an instruction inside the module. The runtime layer's only contribution is the device assignment it built at compile time, plus the coordinator that let a multi-host plugin exchange topology in the first place.

There is no kernel call either, and no stream. Whatever the backend does with warps32 GPU threads scheduled as one unit; the GPU hides latency by switching among resident warps rather than by pipelining a scratchpad.taught in /l/tpu → or vector units or systolic arrays happens entirely behind Compile and Execute. That is the same boundary the xla path opens at /xla/pjrt/pjrt-boundary from the plugin author's side, and the same one chapter 14's walks step through implementation by implementation, the third of them through the very client this lesson maps. This lesson is the view from directly above it.

One leak is worth naming rather than glossing. GetPjRtBuffer sits on the supposedly runtime-neutral interface and returns a PJRT type, because DLPack export needs a raw buffer pointer. The IFRT client throws on it. An abstraction with exactly one working implementation tends to grow a hole shaped like that implementation.