What does PjRtComputationClient actually do?

It is the one implementation of torch_xla’s runtime interface that a normal process ever uses, and it is the only place in the tree that asks PJRT to compile a program, run an executable, or move bytes between the host and a device. Everything above it (the graph executor, the Python bindings, the tensor utilities) holds an abstract ComputationClient pointer and calls virtual methods on it. This file turns each of those calls into one call on an xla::PjRtClient or on a buffer or executable that client handed back. It also owns the small amount of state the seam needs: the PJRT client itself, a map from device id to a dense global ordinal, a map from a device string like TPU:0 back to the PjRtDevice, an operation manager that tracks in-flight work per device, and one hash of the compilation environment computed at startup.

Does the PJRT client cache compiled programs?

No, and the file is unambiguous about it: the word cache appears exactly once in 1,073 lines, inside an error message about the persistent cache, and there is no lookup of any kind before compiling. Compile builds compile options, calls CompileAndLoad on the PJRT client, and returns. The cache lives one layer up, in the graph executor: it hashes the graph plus the compilation environment plus the torch and torch_xla git revisions, looks the hash up in an LRU keyed by that hash, and only calls into this file on a miss. That split is why the counters read the way they do. CachedCompile and UncachedCompile are stamped by the executor and count lookups; CompileTime is stamped in this file and only ever advances when a real compile happens.

What is the difference between ExecuteComputation and ExecuteReplicated?

ExecuteComputation runs one program on one device with one buffer per argument, and calls the PJRT executable’s ExecuteSharded. ExecuteReplicated runs the same program on every local device at once with one shard per device per argument, and calls Execute, whose argument array is indexed device-major. Three things differ beyond the shape of the call. The replicated path transposes its arguments in a thread pool before it can call anything, because torch_xla stores shards per tensor and PJRT wants them per device. It sets strict_shape_checking to true where the single-device path sets it to false. And it takes exactly one lock, on the virtual device named SPMD:0, instead of one lock per device, with a comment in the source saying there is no point grabbing a lock for every individual device on this path.

What happens when I move a tensor to an XLA device?

TransferToDevice is called with a list of tensor sources, and for each one it resolves the device string to a PjRtDevice, adds the tensor’s byte size to a running total, and calls BufferFromHostBuffer with the raw pointer, the element type, the dimensions, the byte strides, and a host-buffer semantics value of kImmutableUntilTransferCompletes. That value is a promise: the runtime may keep reading the host memory after the call returns, and the caller must not free or mutate it until the runtime says it is done. torch_xla keeps that promise with a lambda that captures the tensor shared pointer and has an empty body. The lambda does nothing when it runs; holding a reference until it is destroyed is its entire job. Afterwards the total byte count goes into the OutboundData metric and the number of new handles into a counter.

Why does reading a tensor back from a TPU block?

Because TransferFromDevice waits, on purpose. For each handle it allocates a host literal sized from the buffer’s own shape, starts ToLiteral to fill it, collects the resulting futures, and then awaits all of them before returning. Nothing is returned until every byte has landed. The interface header carries a warning about this: the caller must not hold Python’s global interpreter lock across the call, because the transfer can require other threads to make progress and a held GIL can deadlock the process rather than just slow it down. The practical consequence in a notebook is that any line which needs a real number (printing a loss, calling .item(), calling .cpu()) goes through here and stops your program until the device is finished.

Where torch_xla Calls PJRT: pjrt_computation_client.cpp, Line by Line

Everything PyTorch/XLA ever asks a device to do goes through one 1,073-line C++ file, and this walk reads all of it. Initialize, which runs once per process and quietly decides the names your devices will answer to. TransferToDevice, where a host tensor becomes a device buffer and an empty lambda is the only thing keeping the source alive. Compile, whose 143 lines are almost entirely two arms of an if: one that sets num_partitions to the device count and one that sets num_replicas to it, with a device assignment matrix transposed between them. ExecuteComputation and ExecuteReplicated, which take the same executable, disagree about strict shape checking, and lock in completely different ways. TransferFromDevice, the one call that genuinely blocks. And a first function in the file that is never called at all. The signature exhibit runs one step through the file in its own order, each crossing showing the PJRT method underneath it and the timer it stamps.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of Where torch_xla Calls PJRT: pjrt_computation_client.cpp, Line by Line — built to be understood, not skimmed.

Questions

What does PjRtComputationClient actually do?: It is the one implementation of torch_xla’s runtime interface that a normal process ever uses, and it is the only place in the tree that asks PJRT to compile a program, run an executable, or move bytes between the host and a device. Everything above it (the graph executor, the Python bindings, the tensor utilities) holds an abstract ComputationClient pointer and calls virtual methods on it. This file turns each of those calls into one call on an xla::PjRtClient or on a buffer or executable that client handed back. It also owns the small amount of state the seam needs: the PJRT client itself, a map from device id to a dense global ordinal, a map from a device string like TPU:0 back to the PjRtDevice, an operation manager that tracks in-flight work per device, and one hash of the compilation environment computed at startup.
Does the PJRT client cache compiled programs?: No, and the file is unambiguous about it: the word cache appears exactly once in 1,073 lines, inside an error message about the persistent cache, and there is no lookup of any kind before compiling. Compile builds compile options, calls CompileAndLoad on the PJRT client, and returns. The cache lives one layer up, in the graph executor: it hashes the graph plus the compilation environment plus the torch and torch_xla git revisions, looks the hash up in an LRU keyed by that hash, and only calls into this file on a miss. That split is why the counters read the way they do. CachedCompile and UncachedCompile are stamped by the executor and count lookups; CompileTime is stamped in this file and only ever advances when a real compile happens.
What is the difference between ExecuteComputation and ExecuteReplicated?: ExecuteComputation runs one program on one device with one buffer per argument, and calls the PJRT executable’s ExecuteSharded. ExecuteReplicated runs the same program on every local device at once with one shard per device per argument, and calls Execute, whose argument array is indexed device-major. Three things differ beyond the shape of the call. The replicated path transposes its arguments in a thread pool before it can call anything, because torch_xla stores shards per tensor and PJRT wants them per device. It sets strict_shape_checking to true where the single-device path sets it to false. And it takes exactly one lock, on the virtual device named SPMD:0, instead of one lock per device, with a comment in the source saying there is no point grabbing a lock for every individual device on this path.
What happens when I move a tensor to an XLA device?: TransferToDevice is called with a list of tensor sources, and for each one it resolves the device string to a PjRtDevice, adds the tensor’s byte size to a running total, and calls BufferFromHostBuffer with the raw pointer, the element type, the dimensions, the byte strides, and a host-buffer semantics value of kImmutableUntilTransferCompletes. That value is a promise: the runtime may keep reading the host memory after the call returns, and the caller must not free or mutate it until the runtime says it is done. torch_xla keeps that promise with a lambda that captures the tensor shared pointer and has an empty body. The lambda does nothing when it runs; holding a reference until it is destroyed is its entire job. Afterwards the total byte count goes into the OutboundData metric and the number of new handles into a counter.
Why does reading a tensor back from a TPU block?: Because TransferFromDevice waits, on purpose. For each handle it allocates a host literal sized from the buffer’s own shape, starts ToLiteral to fill it, collects the resulting futures, and then awaits all of them before returning. Nothing is returned until every byte has landed. The interface header carries a warning about this: the caller must not hold Python’s global interpreter lock across the call, because the transfer can require other threads to make progress and a held GIL can deadlock the process rather than just slow it down. The practical consequence in a notebook is that any line which needs a real number (printing a loss, calling .item(), calling .cpu()) goes through here and stops your program until the device is finished.

Where torch_xla Calls PJRT: pjrt_computation_client.cpp, Line by Line

Questions

Related explainers