PJRT is OpenXLA’s uniform device API for connecting machine learning frameworks to hardware-specific compiler and runtime implementations. It sits below frameworks such as JAX, TensorFlow, and PyTorch/XLA and above concrete devices such as TPU, GPU, CPU, or custom accelerators. Instead of every framework integrating separately with every accelerator backend, a hardware vendor can implement a PJRT plugin once and expose that implementation through a stable C API.
The problem PJRT solves is integration combinatorics. Without a common device boundary, each framework, compiler, and runtime pair tends to become a bespoke vertical stack. That slows down new accelerator adoption, increases maintenance cost, and makes workload portability fragile. PJRT turns that matrix into a smaller set of framework-to-PJRT and PJRT-to-device contracts.
PJRT complements, rather than replaces, XLA and StableHLO. StableHLO is the portable operation-level representation that lets frameworks describe computations. XLA, IREE, or another compiler can lower that representation. PJRT is the runtime interface that creates clients, manages devices and buffers, compiles or loads executables, and dispatches execution on the selected backend.
A PJRT client encapsulates the state and communication path for a backend. Framework code calls generic PJRT operations, while the plugin hides device-specific details such as command queues, compiler invocation, memory layout, and synchronization. This opacity is what lets a framework load multiple backend plugins without embedding each backend’s implementation strategy.
The plugin boundary is usually a dynamically loaded library that exports a PJRT C API table. A framework discovers or imports the plugin, calls initialization hooks, creates a client, transfers buffers, compiles a computation, and invokes executable entry points. The binary can be open source or closed source as long as it obeys the PJRT ABI and version contract.
PJRT models execution in terms of devices, memory spaces, buffers, and loaded executables. Buffers represent data owned by a backend; devices describe where computation can run; memory spaces describe where data may live; loaded executables are compiled artifacts ready to execute with PJRT buffers as arguments.
XLA decides how a computation should be compiled. PJRT defines how a framework talks to the backend that owns compilation, memory transfer, executable loading, execution, and readiness signaling.
A plugin exposes a function that returns a table of function pointers. The most important early hook is client creation: the framework asks the plugin for a client, and the plugin returns an object that can enumerate devices, create buffers, compile modules, and run executables.
#include "xla/pjrt/c/pjrt_c_api.h"
#include "xla/pjrt/c/pjrt_c_api_wrapper_impl.h"
namespace my_plugin {
PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) {
std::unique_ptr<xla::PjRtClient> client = GetPluginPjRtClient();
args->client = pjrt::CreateWrapperClient(std::move(client));
return nullptr;
}
const PJRT_Api* GetPjrtApi() {
static const PJRT_Api api =
pjrt::CreatePjrtApi(my_plugin::PJRT_Client_Create);
return &api;
}
} // namespace my_pluginAfter the framework has a client, it can transfer inputs into backend-owned buffers, compile a module into an executable, and execute that artifact on one or more devices. The backend decides whether compilation is delegated to XLA, a proprietary compiler, an MLIR pipeline, or a reference interpreter.
Device execution is frequently asynchronous. PJRT buffers expose readiness futures so a framework can enqueue work, continue running Python or host-side logic, and block only when the value is actually needed. A correct plugin must make these readiness signals reflect real device completion, not just host-side enqueue success.
A hardware team building an accelerator can implement a PJRT plugin to make JAX or PyTorch/XLA workloads target the device without writing a full framework fork. The plugin packages the compiler, runtime hooks, device discovery, and memory movement behind one interface.
Because PJRT is used by JAX, TensorFlow paths, and PyTorch/XLA, a backend implementation can serve multiple framework ecosystems. This is especially valuable for vendors that need broad developer coverage before their accelerator has a mature native framework stack.
PJRT is also useful for experimental compiler stacks. A team can expose a StableHLO interpreter, an MLIR-based backend, or a simulator as a PJRT device, then validate framework integration and semantic correctness before optimizing device kernels.
The common implementation path is to write C++ classes for the client, device, memory space, buffer, executable, and compiler behavior, then expose them through the PJRT C API wrapper. This preserves a stable plugin ABI while letting the backend use richer C++ abstractions internally.
class MyPjRtClient : public xla::PjRtClient {
public:
absl::string_view platform_name() const override { return "my_accel"; }
absl::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable>> Compile(
const xla::CompileOptions& options,
std::unique_ptr<xla::PjRtProgram> program) override {
auto artifact = compiler_.LowerAndCodegen(*program, options);
return std::make_unique<MyLoadedExecutable>(std::move(artifact), runtime_);
}
absl::StatusOr<std::unique_ptr<xla::PjRtBuffer>> BufferFromHostBuffer(
const void* data, xla::PrimitiveType type,
absl::Span<const int64_t> dims,
std::optional<absl::Span<const int64_t>> byte_strides,
xla::PjRtClient::HostBufferSemantics semantics,
std::function<void()> on_done,
xla::PjRtMemorySpace* memory_space,
const xla::Layout* device_layout) override {
return runtime_.CopyToDevice(data, type, dims, semantics, on_done);
}
};Use StableHLO as the portable computation format whenever possible. That keeps the frontend independent from backend compiler details and makes it easier to compare plugin behavior against reference implementations.
Unit tests for the C API are necessary but insufficient. Production plugins should run framework-level tests that exercise JIT compilation, buffer lifetimes, asynchronous reads, shape polymorphism, sharding, and error propagation from the user’s point of view.