What is ComputationClient in PyTorch/XLA?

It is the abstract C++ class that every crossing out of PyTorch and into the XLA runtime goes through. It declares 43 pure-virtual methods: put a tensor on a device, compile a computation, execute it on one device or on all of them, read a result back, and answer questions about the devices you have. It holds no state of its own beyond a few static metric handles, so it is a contract, not a component. One implementation ships and runs, PjRtComputationClient; a second, IfrtComputationClient, is compiled but switched off.

Why are there four different things called Computation in torch_xla?

Because four layers each need their own word for the same idea and nobody renamed them. The header says so directly in a comment above the class. torch::lazy::Computation is the lazy-tensor core’s general notion of a computation. ComputationClient::Computation is torch_xla’s wrapper: one xla::XlaComputation plus the list of devices it targets. xla::XlaComputation is the compiler’s own object, the HLO program. And PjRtComputationClient::PjRtComputation inherits from the wrapper and adds the handle to the compiled executable. The wrapper is reused for three different purposes, which the header also spells out and calls "not ideal".

Why does TransferFromDevice warn about the GIL?

Because it blocks. The comment above it says the call will wait until the buffers behind the handles are ready if they came from TransferToDevice or an Execute call, and that calling it from Python while holding the GIL can deadlock. The device work that has to finish may itself need to run Python (a callback, a host transfer, a custom call), and it cannot, because the thread that would run it is parked inside the transfer. Reading a value back is the one place in the interface where a Python-side habit turns into a hang.

What does HashCompilationEnv do and why does the graph executor call it?

It returns a hash of everything about the machine and the build that could change what the compiler emits: the XLA environment variables, and either the serialized topology description of the client or, failing that, the platform name, the platform version (which carries the libtpu version and the hardware type), and every global device string in a fixed order. The graph executor folds that hash into the key for every graph it compiles, alongside the Torch and XLA git revisions. Without it, a persistent compilation cache written by yesterday’s libtpu on a different chip generation would happily be reused today.

Does PyTorch/XLA use IFRT?

Not at this commit. Both clients are compiled into the library, but the switch that would pick IFRT is a hard-coded constant set to false in runtime.cpp, with the environment-variable read that used to drive it commented out just above and a TODO saying to enable IFRT once it stops crashing. Even if you flipped it, IFRT’s ExecuteComputation is a stub that returns an unimplemented error. So every crossing in a real run lands on PjRtComputationClient.

The Seam Interface: computation_client.h, Line by Line

Everything PyTorch does on an XLA device leaves the process through one C++ class. Walk the real computation_client.h top to bottom: the four different things called "Computation", the transfer surface and its GIL warning, Compile, ExecuteReplicated, and the 43 pure-virtual methods that are the entire seam between PyTorch and PJRT.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The Seam Interface: computation_client.h, Line by Line — built to be understood, not skimmed.

Questions

What is ComputationClient in PyTorch/XLA?: It is the abstract C++ class that every crossing out of PyTorch and into the XLA runtime goes through. It declares 43 pure-virtual methods: put a tensor on a device, compile a computation, execute it on one device or on all of them, read a result back, and answer questions about the devices you have. It holds no state of its own beyond a few static metric handles, so it is a contract, not a component. One implementation ships and runs, PjRtComputationClient; a second, IfrtComputationClient, is compiled but switched off.
Why are there four different things called Computation in torch_xla?: Because four layers each need their own word for the same idea and nobody renamed them. The header says so directly in a comment above the class. torch::lazy::Computation is the lazy-tensor core’s general notion of a computation. ComputationClient::Computation is torch_xla’s wrapper: one xla::XlaComputation plus the list of devices it targets. xla::XlaComputation is the compiler’s own object, the HLO program. And PjRtComputationClient::PjRtComputation inherits from the wrapper and adds the handle to the compiled executable. The wrapper is reused for three different purposes, which the header also spells out and calls "not ideal".
Why does TransferFromDevice warn about the GIL?: Because it blocks. The comment above it says the call will wait until the buffers behind the handles are ready if they came from TransferToDevice or an Execute call, and that calling it from Python while holding the GIL can deadlock. The device work that has to finish may itself need to run Python (a callback, a host transfer, a custom call), and it cannot, because the thread that would run it is parked inside the transfer. Reading a value back is the one place in the interface where a Python-side habit turns into a hang.
What does HashCompilationEnv do and why does the graph executor call it?: It returns a hash of everything about the machine and the build that could change what the compiler emits: the XLA environment variables, and either the serialized topology description of the client or, failing that, the platform name, the platform version (which carries the libtpu version and the hardware type), and every global device string in a fixed order. The graph executor folds that hash into the key for every graph it compiles, alongside the Torch and XLA git revisions. Without it, a persistent compilation cache written by yesterday’s libtpu on a different chip generation would happily be reused today.
Does PyTorch/XLA use IFRT?: Not at this commit. Both clients are compiled into the library, but the switch that would pick IFRT is a hard-coded constant set to false in runtime.cpp, with the environment-variable read that used to drive it commented out just above and a TODO saying to enable IFRT once it stops crashing. Even if you flipped it, IFRT’s ExecuteComputation is a stub that returns an unimplemented error. So every crossing in a real run lands on PjRtComputationClient.

The Seam Interface: computation_client.h, Line by Line

Questions

Related explainers