What is PJRT in PyTorch/XLA?

PJRT is the runtime interface XLA exposes to whoever wants to run a compiled program on an accelerator. It is a small C++ interface, with a matching C ABI so hardware vendors can ship plugins, and it answers roughly four kinds of request: put these host bytes on the device, compile this program, run that program, give me these bytes back. PyTorch/XLA never calls it directly from its lazy-tensor code. Everything funnels through one pure-virtual class of its own, ComputationClient (declared at torch_xla/csrc/runtime/computation_client.h:55), and exactly one implementation of that class, PjRtComputationClient, is the only code in torch_xla that ever touches a PJRT method. That indirection is why a second runtime (IFRT) can be compiled into the same binary and switched off without anything above the line noticing.

How many times does a torch_xla training step call into PJRT?

Four, in lazy mode, and always in the same order. TransferToDevice sends your parameters and batch down, landing on PJRT’s BufferFromHostBuffer. Compile hands the lowered HLO program down, landing on CompileAndLoad, which both compiles it and loads it onto the devices. ExecuteComputation runs the loaded executable, landing on ExecuteSharded, and returns immediately with buffers that are not ready yet. Then TransferFromDevice pulls one value back with ToLiteral, and that call blocks. A published capture of one small MLP step on a Colab TPU v6e-1 with torch_xla 2.9.0 reported exactly one sample on each of the four corresponding timers. The dozen operations in the step never reach the runtime individually; they are recorded, fused into one program, and carried down once. Eager mode is the opposite trade: the same capture, run with the barrier after every operation, recorded 23 compiles and 58 executes over two steps.

Does an all-reduce cross the PJRT boundary at run time?

No, and there is no method on the interface that could carry it. Collectives are lowered into the program: by the time Compile runs, the all-reduce is already an instruction in the HLO, sitting between a matmul and a multiply like any other operation. The compiler plans it and the executable performs it. The only place collective code touches the runtime at all is one query in cross_replica_reduces.cpp asking how many devices there are so it can size a split, which reads configuration and moves no data. Buffer donation works the same way: it is expressed as buffer-donor annotations attached to the HLO during lowering (xla_graph_executor.cpp:1431), never as a runtime option. You can see the absence in the type system, because the options struct both execute methods take holds exactly two booleans, one for unpacking a tuple result and one for picking a metric name. The practical consequence: if you are chasing a slow all-reduce, the seam counters will tell you nothing and the compiler dump will tell you everything.

Why does torch_xla need PJRT_DEVICE to be set?

Because the runtime refuses to build a client without knowing which one to build. The check happens before anything else, at torch_xla/csrc/runtime/runtime.cpp:34, and the message is literally “$PJRT_DEVICE is not set.” It fires before the registry that maps a device name to a client is ever consulted, which is worth knowing: the error means the client was never constructed, so nothing about your model or your tensors is wrong yet. Set it to CPU and you get an in-process client with no shared library involved at all, which is the right way to learn this stack on a laptop. Set it to TPU and torch_xla dlopens libtpu, initializes it as a PJRT plugin, and asks for a C-API client, which is the same three-step dance any other plugin device goes through. The client is a per-process singleton, so a second attempt to build one in the same process is a hard failure with the message “ComputationClient can only be initialized once.”

Can torch_xla use IFRT instead of PJRT?

Not at this commit, and not by setting anything. IfrtComputationClient exists in the tree, is about seven hundred lines, and compiles. It is switched off by a hard-coded const bool use_ifrt = false at torch_xla/csrc/runtime/runtime.cpp:31, directly beneath a commented-out environment read and a TODO saying to enable it once it stops crashing. Flipping the boolean would not get you far: that client only accepts SPMD compilations, and its ExecuteComputation returns “ExecuteComputation not implemented”. The interesting part is what the arrangement demonstrates rather than what it delivers. A whole second runtime implementation can sit compiled inside the binary, share every metric name (the counters are static members of the interface, not the implementation), and be disabled with one line, precisely because everything above the seam talks to the abstract class and never to a concrete one.

The PJRT Boundary: One Training Step, Crossing by Crossing

Somewhere between your PyTorch code and a TPU there is a line, and it is one C++ class: ComputationClient, pure virtual, forty-three methods, one live implementation. This page walks a single lazy training step across it four times (inputs down, program down, run, one value up), names the PJRT method waiting on the other side of each crossing, and counts them against a published lab capture on a Colab TPU. Then the things that never cross: an all-reduce is an instruction inside the program, not a call, and buffer donation rides down as an HLO annotation because the execute options struct is two booleans wide. Closes with the metrics report that has the seam running down its middle, the fourteen files that reach the client, the SPMD path that swaps one method for its plural, and the pinned OpenXLA commit that makes every line number checkable.

Concept · AI / ML. The source ↗

A free, interactive, animated visual explainer of The PJRT Boundary: One Training Step, Crossing by Crossing — built to be understood, not skimmed.

Questions

What is PJRT in PyTorch/XLA?: PJRT is the runtime interface XLA exposes to whoever wants to run a compiled program on an accelerator. It is a small C++ interface, with a matching C ABI so hardware vendors can ship plugins, and it answers roughly four kinds of request: put these host bytes on the device, compile this program, run that program, give me these bytes back. PyTorch/XLA never calls it directly from its lazy-tensor code. Everything funnels through one pure-virtual class of its own, ComputationClient (declared at torch_xla/csrc/runtime/computation_client.h:55), and exactly one implementation of that class, PjRtComputationClient, is the only code in torch_xla that ever touches a PJRT method. That indirection is why a second runtime (IFRT) can be compiled into the same binary and switched off without anything above the line noticing.
How many times does a torch_xla training step call into PJRT?: Four, in lazy mode, and always in the same order. TransferToDevice sends your parameters and batch down, landing on PJRT’s BufferFromHostBuffer. Compile hands the lowered HLO program down, landing on CompileAndLoad, which both compiles it and loads it onto the devices. ExecuteComputation runs the loaded executable, landing on ExecuteSharded, and returns immediately with buffers that are not ready yet. Then TransferFromDevice pulls one value back with ToLiteral, and that call blocks. A published capture of one small MLP step on a Colab TPU v6e-1 with torch_xla 2.9.0 reported exactly one sample on each of the four corresponding timers. The dozen operations in the step never reach the runtime individually; they are recorded, fused into one program, and carried down once. Eager mode is the opposite trade: the same capture, run with the barrier after every operation, recorded 23 compiles and 58 executes over two steps.
Does an all-reduce cross the PJRT boundary at run time?: No, and there is no method on the interface that could carry it. Collectives are lowered into the program: by the time Compile runs, the all-reduce is already an instruction in the HLO, sitting between a matmul and a multiply like any other operation. The compiler plans it and the executable performs it. The only place collective code touches the runtime at all is one query in cross_replica_reduces.cpp asking how many devices there are so it can size a split, which reads configuration and moves no data. Buffer donation works the same way: it is expressed as buffer-donor annotations attached to the HLO during lowering (xla_graph_executor.cpp:1431), never as a runtime option. You can see the absence in the type system, because the options struct both execute methods take holds exactly two booleans, one for unpacking a tuple result and one for picking a metric name. The practical consequence: if you are chasing a slow all-reduce, the seam counters will tell you nothing and the compiler dump will tell you everything.
Why does torch_xla need PJRT_DEVICE to be set?: Because the runtime refuses to build a client without knowing which one to build. The check happens before anything else, at torch_xla/csrc/runtime/runtime.cpp:34, and the message is literally “$PJRT_DEVICE is not set.” It fires before the registry that maps a device name to a client is ever consulted, which is worth knowing: the error means the client was never constructed, so nothing about your model or your tensors is wrong yet. Set it to CPU and you get an in-process client with no shared library involved at all, which is the right way to learn this stack on a laptop. Set it to TPU and torch_xla dlopens libtpu, initializes it as a PJRT plugin, and asks for a C-API client, which is the same three-step dance any other plugin device goes through. The client is a per-process singleton, so a second attempt to build one in the same process is a hard failure with the message “ComputationClient can only be initialized once.”
Can torch_xla use IFRT instead of PJRT?: Not at this commit, and not by setting anything. IfrtComputationClient exists in the tree, is about seven hundred lines, and compiles. It is switched off by a hard-coded const bool use_ifrt = false at torch_xla/csrc/runtime/runtime.cpp:31, directly beneath a commented-out environment read and a TODO saying to enable it once it stops crashing. Flipping the boolean would not get you far: that client only accepts SPMD compilations, and its ExecuteComputation returns “ExecuteComputation not implemented”. The interesting part is what the arrangement demonstrates rather than what it delivers. A whole second runtime implementation can sit compiled inside the binary, share every metric name (the counters are static members of the interface, not the implementation), and be disabled with one line, precisely because everything above the seam talks to the abstract class and never to a concrete one.

The PJRT Boundary: One Training Step, Crossing by Crossing

Questions

Related explainers