PJRT, the boundary

The chapter above this lesson opens at the seam and stops there, which is the right altitude for orientation: it names the boundary, names the three nouns on either side of it, and moves on. This lesson starts one line below that, in the symbol table of a plugin's shared object. Read the linker version script XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → ships for its own CPU plugin and the entire public surface of a PJRT backend turns out to be a single name. GetPjrtApi is global. Everything else in the library is local.

The spelling is worth pinning down before you type it from memory, because the tree and its own documentation disagree. Every header in xla/pjrt/c declares the function as GetPjrtApi, lowercase r and t, and both version scripts export it under that spelling. The integration guide's prose calls the step GetPjRtApi. Match the header.

The declaration is three lines, and the comment above it carries an ownership rule that costs you a crash if you miss it. The caller does not take ownership of what comes back, so the struct a plugin returns has to outlive every call made through it. In practice that means a function-local static, built once, returned by pointer forever after.

#include "xla/pjrt/c/pjrt_c_api.h"

#ifdef __cplusplus
extern "C" {
#endif

// Does not pass ownership of returned PJRT_Api* to caller.
const PJRT_Api* GetPjrtApi();

#ifdef __cplusplus
}
#endif

PJRT_Api is a struct of function pointers, one member per operation, and there are two honest ways to fill it in. The first route implements the C API directly. You write a PJRT_Error* function for every member, you define your own concrete types behind the opaque PJRT_Client and PJRT_Buffer handles, and you hand the finished struct back. Nothing in XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → needs to know your project exists, and nothing in your project needs to link XLA.

The second route is shorter by a few thousand lines, and it is the one the tree's own example plugin takes. Write an ordinary C++ subclass of xla::PjRtClient, the same class an in-tree backend would write, then let XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → generate the C surface around it. pjrt::CreateWrapperClient takes that C++ client and returns the opaque PJRT_Client handle the framework will hold. pjrt::CreatePjrtApi builds the whole PJRT_Api out of XLA's stock implementations, and asks you for only the pieces it cannot invent.

Look at what CreatePjrtApi's parameter list actually demands and the size of a minimal plugin becomes concrete. Four creates: a client, an execute context, a topology description, and a plugin initializer. The last two parameters have defaults, and the example plugin fills the initializer with pjrt::PJRT_Plugin_Initialize_NoOp because it has nothing to set up. Everything else in the struct, every buffer call and every executable call, is XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla →'s wrapper code forwarding into your C++ client's virtual methods.

// xla/pjrt/c/pjrt_c_api_wrapper_impl.h
PJRT_Client* CreateWrapperClient(const PJRT_Api* api,
                                 std::unique_ptr<xla::PjRtClient> cpp_client);

PJRT_Api CreatePjrtApi(
    PJRT_Client_Create* create_fn,
    PJRT_ExecuteContext_Create* execute_context_create_fn,
    PJRT_TopologyDescription_Create* topology_create_fn,
    PJRT_Plugin_Initialize* plugin_initialize_fn,
    PJRT_Extension_Base* extension_start = nullptr,
    PJRT_Plugin_Attributes* plugin_attributes_fn =
        pjrt::PJRT_Plugin_Attributes_Empty);

// xla/pjrt/plugin/example_plugin/myplugin_c_pjrt_internal.cc
PJRT_Error* PJRT_MypluginClient_Create(PJRT_Client_Create_Args* args) {
  std::unique_ptr<xla::PjRtClient> client = CreateMyPluginPjrtClient();
  args->client =
      pjrt::CreateWrapperClient(GetMyPluginPjrtApi(), std::move(client));
  return nullptr;
}

static const PJRT_Api pjrt_api = pjrt::CreatePjrtApi(
    myplugin_pjrt::PJRT_MypluginClient_Create,
    myplugin_pjrt::PJRT_MypluginExecuteContext_Create,
    myplugin_pjrt::PJRT_MypluginDeviceTopology_Create,
    pjrt::PJRT_Plugin_Initialize_NoOp, &example_extension.base,
    pjrt::PJRT_Plugin_Attributes_Xla);

Every function in the C API takes exactly one argument, a pointer to its own Args struct, and returns PJRT_Error*, where a null pointer means success. The typedef reads typedef PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args); and the pattern repeats for all of them. Inputs and outputs both live in the struct, so reading an Args declaration tells you the whole contract of a call without hunting for a return type.

PJRT_Client_Create_Args is worth reading field by field, because three of its members carry a distributed system into a struct that otherwise looks single-host. The kv_get, kv_put, and kv_try_get callbacks are the framework handing your plugin a key-value store it already owns. When a multi-host job starts, that store is how a plugin on host 3 learns what a plugin on host 0 decided, and it is the concrete thing on the far side of the jax.distributed.initialize() call the path meets at /xla/mcjax.

One detail in the field ordering is a small lesson in reading C ABIs. The output field, client, sits in the middle of the struct rather than at the end, with the kv_try_get pair after it. Fields get appended, never inserted, because inserting one would move every field after it and break every already-compiled caller. So the tail of any Args struct is a rough history of what got added last.

A shared object with the right symbol in it is inert until something loads it, and JAX has two ways of finding one. The first is a namespace package: create a directory called jax_plugins, put a module under it, and JAX will import it. The second is packaging metadata, an entry point registered under the group jax_plugins in your pyproject.toml. discover_pjrt_plugins() in jax/_src/xla_bridge.py walks both, and for each thing it finds it calls one function by name, initialize().

What initialize() has to do is call xb.register_plugin. The signature is worth copying exactly rather than approximating, because two of its arguments are alternatives rather than companions: library_path points at the .so JAX should dlopen, and c_api is for a plugin already resident in the process. The default priority is 400, so a plugin that wants to outrank the stock backends passes a higher number.

Priority only decides anything when JAX_PLATFORMS is unset. Set that variable and selection becomes explicit, and it fails loudly rather than quietly falling back to CPU, which is the behaviour you want while a new plugin is still half-written. A plugin that silently loses to CPU looks exactly like a plugin that loaded fine and produced wrong devices.

# jax_plugins/my_backend/__init__.py
from jax._src import xla_bridge as xb


def initialize() -> None:
    xb.register_plugin(
        "my_backend",
        priority=500,          # any value above the 400 default outranks cpu
        library_path="/opt/my_backend/libmy_backend_pjrt_plugin.so",
        options=None,
    )


# the signature you are calling, from jax/_src/xla_bridge.py:
#
# def register_plugin(plugin_name, *, priority=400, library_path=None,
#                     options=None, c_api=None, factory=None,
#                     make_topology=None)

The header carries its own version in two macros and one struct. PJRT_API_MAJOR is 0 and PJRT_API_MINOR is 114 at the commit this lesson was written against. The major number is incremented when an ABI-incompatible change is made to the interface; the minor is incremented when the interface is updated in a way that is potentially ABI-compatible with older versions. PJRT_Api_Version carries both across the boundary at runtime, so a framework can ask a plugin what it was built against before calling anything newer than that.

That machinery describes an aspiration more than the current rule. The integration guide states the practical constraint plainly: you need to match the jaxlib version with the PJRT C API version, and the recommendation for out-of-tree work is nightly jaxlib from the same date as 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 → commit you built against. ABI stability is described as coming, not arrived. Budget for rebuilding your plugin against each jaxlib you support, and treat struct_size and extension_start as the mechanism that makes that rebuild cheap rather than one that makes it unnecessary.

libtpu is this scheme running in production, which is the useful thing to hold in mind whenever the TPU stack feels like a special case. It is a wheel on PyPI, described by its own publisher as the Google Cloud TPU runtime library, providing core functionality for compilation, inter-chip communication, and runtime execution. Version 0.0.45 went out on 2026-08-01 at 215.9 MB per wheel. JAX finds it through the same plugin discovery as anything else, dlopens it, and calls the same three-line GetPjrtApi that the CPU header declares.

The four-line header the CPU plugin exports and the 215 MB TPU wheel present the identical symbol.

So the difference between a toy plugin and libtpu is not the interface. It is what sits behind the function pointers, which is a compiler and a runtime for a chip nobody outside Google can read. The codegen unit's lesson at /xla/codegen/backend-seam takes that seam apart properly.