What does PJRT_DEVICE actually do in PyTorch/XLA?

It is a plain environment variable holding a device-type string, and it is read in exactly two places. The first is runtime.cpp, which refuses to build a runtime at all if the variable is empty and returns the error "$PJRT_DEVICE is not set." The second is pjrt_registry.cpp, where the string is compared against a fixed if/else chain to decide which PJRT client to construct: an in-process CPU client, a TPU client loaded from libtpu.so, an XPU or NEURON client loaded from their vendor libraries, or a client built from a plugin registered at run time. The comparisons are exact string equality, so PJRT_DEVICE=tpu in lowercase does not match the TPU branch and falls all the way through to the unknown-device error. There is no partial matching, no aliasing, and no default: an unrecognised string is an InvalidArgumentError naming the value you passed.

What is a PJRT plugin?

A shared library that implements the PJRT C API, plus a small Python object that tells PyTorch/XLA where to find it. On the C++ side the plugin object answers three questions and nothing else: where its library lives on disk (library_path), what options to hand the client constructor (client_create_options), and whether it needs a distributed key-value store before the client can be built (requires_xla_coordinator). Loading it is three calls into OpenXLA: LoadPjrtPlugin dlopens the library and looks up a single exported symbol named GetPjrtApi, InitializePjrtPlugin lets the plugin set itself up, and GetCApiClient wraps the resulting function table in an xla::PjRtClient the rest of torch_xla can use. Because the boundary is a C function table rather than C++ headers, a vendor can ship a plugin built against a different compiler and standard library than the framework loading it.

How does PyTorch/XLA find libtpu.so?

Three places, in order, and the order is the point. First TPU_LIBRARY_PATH, which is the variable you set by hand and which nothing in torch_xla ever overwrites. Second PTXLA_TPU_LIBRARY_PATH, which torch_xla sets for itself during import when it can infer a path: either the libtpu.so bundled inside the installed torch_xla package, or the one the libtpu pip wheel reports. Third the bare string "libtpu.so", which hands the problem to the dynamic linker and its search path. The two-variable split exists so the inference does not collide with another framework in the same process: JAX also reads TPU_LIBRARY_PATH, so torch_xla writes its guess into a private variable instead, and the source comment says the private one will be removed in a future version.

Does PyTorch/XLA still support PJRT_DEVICE=CUDA?

Not through this file. At this commit there is no CUDA branch in the if/else chain and no CUDA entry point among the plugins the package installs, which are tpu, neuron, and xpu. The CUDA-specific runtime logic was deleted upstream in PR #9598, and a plain PJRT_DEVICE=CUDA now falls past every branch to the final else and fails with an InvalidArgumentError reading Unknown PJRT_DEVICE: ‘CUDA’. The one path that still exists is the generic plugin path: if PJRT_DYNAMIC_PLUGINS is on and something has registered a plugin under the name CUDA before the runtime initializes, the first branch takes it and loads that plugin like any other. Nothing in the tree registers one, so that has to come from an out-of-tree package.

What is the XlaCoordinator, and when is it created?

It is a thin wrapper over the OpenXLA distributed runtime that PyTorch/XLA uses for two things: coordinating preemption notices across hosts, and acting as a shared key-value store so the processes on different hosts can find each other while their device clients start up. In pjrt_registry.cpp it is created on exactly one condition, that the plugin returns true from requires_xla_coordinator. When it is, the file reads the rank and world size from PJRT_LOCAL_PROCESS_RANK, LOCAL_RANK, RANK, PJRT_LOCAL_PROCESS_COUNT, LOCAL_WORLD_SIZE and WORLD_SIZE (the torchrun names are the fallbacks), reads the address from MASTER_ADDR and XLA_COORDINATOR_PORT, starts the coordinator, and passes the resulting store into GetCApiClient under the key prefix "pjrt:". None of the plugins shipped in the tree override requires_xla_coordinator, so on a stock install that block is skipped and the coordinator stays null.

How PJRT_DEVICE Becomes a Client: pjrt_registry.cpp, Line by Line

Set PJRT_DEVICE=TPU, import torch_xla, and a chip you never named starts running your model. The whole decision is 174 lines of C++ in one file. We read all of it: the plugin interface with its three questions, the global map seeded with a single placeholder entry, the exact-string lookup whose error message doubles as the documentation, and then the if/else chain itself, branch by branch. The dynamic-plugin branch that a plain import turns on by default, and which dlopens a vendor .so, initializes it, and wraps it through the PJRT C API. The distributed key-value store that only exists for plugins that ask for it, built out of the same coordinator torch_xla uses for preemption. The CPU branch, the one device with no plugin at all. The TPU branch and its three-deep search for libtpu.so. The two half-finished branches, XPU and NEURON, that skip the plugin initialize and the profiler hook. And the else that catches your typo. The signature exhibit runs the real branch order: type a device string, flip the dynamic-plugin switch, and watch which branch catches it and which environment variables it reads.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of How PJRT_DEVICE Becomes a Client: pjrt_registry.cpp, Line by Line — built to be understood, not skimmed.

Questions

What does PJRT_DEVICE actually do in PyTorch/XLA?: It is a plain environment variable holding a device-type string, and it is read in exactly two places. The first is runtime.cpp, which refuses to build a runtime at all if the variable is empty and returns the error "$PJRT_DEVICE is not set." The second is pjrt_registry.cpp, where the string is compared against a fixed if/else chain to decide which PJRT client to construct: an in-process CPU client, a TPU client loaded from libtpu.so, an XPU or NEURON client loaded from their vendor libraries, or a client built from a plugin registered at run time. The comparisons are exact string equality, so PJRT_DEVICE=tpu in lowercase does not match the TPU branch and falls all the way through to the unknown-device error. There is no partial matching, no aliasing, and no default: an unrecognised string is an InvalidArgumentError naming the value you passed.
What is a PJRT plugin?: A shared library that implements the PJRT C API, plus a small Python object that tells PyTorch/XLA where to find it. On the C++ side the plugin object answers three questions and nothing else: where its library lives on disk (library_path), what options to hand the client constructor (client_create_options), and whether it needs a distributed key-value store before the client can be built (requires_xla_coordinator). Loading it is three calls into OpenXLA: LoadPjrtPlugin dlopens the library and looks up a single exported symbol named GetPjrtApi, InitializePjrtPlugin lets the plugin set itself up, and GetCApiClient wraps the resulting function table in an xla::PjRtClient the rest of torch_xla can use. Because the boundary is a C function table rather than C++ headers, a vendor can ship a plugin built against a different compiler and standard library than the framework loading it.
How does PyTorch/XLA find libtpu.so?: Three places, in order, and the order is the point. First TPU_LIBRARY_PATH, which is the variable you set by hand and which nothing in torch_xla ever overwrites. Second PTXLA_TPU_LIBRARY_PATH, which torch_xla sets for itself during import when it can infer a path: either the libtpu.so bundled inside the installed torch_xla package, or the one the libtpu pip wheel reports. Third the bare string "libtpu.so", which hands the problem to the dynamic linker and its search path. The two-variable split exists so the inference does not collide with another framework in the same process: JAX also reads TPU_LIBRARY_PATH, so torch_xla writes its guess into a private variable instead, and the source comment says the private one will be removed in a future version.
Does PyTorch/XLA still support PJRT_DEVICE=CUDA?: Not through this file. At this commit there is no CUDA branch in the if/else chain and no CUDA entry point among the plugins the package installs, which are tpu, neuron, and xpu. The CUDA-specific runtime logic was deleted upstream in PR #9598, and a plain PJRT_DEVICE=CUDA now falls past every branch to the final else and fails with an InvalidArgumentError reading Unknown PJRT_DEVICE: ‘CUDA’. The one path that still exists is the generic plugin path: if PJRT_DYNAMIC_PLUGINS is on and something has registered a plugin under the name CUDA before the runtime initializes, the first branch takes it and loads that plugin like any other. Nothing in the tree registers one, so that has to come from an out-of-tree package.
What is the XlaCoordinator, and when is it created?: It is a thin wrapper over the OpenXLA distributed runtime that PyTorch/XLA uses for two things: coordinating preemption notices across hosts, and acting as a shared key-value store so the processes on different hosts can find each other while their device clients start up. In pjrt_registry.cpp it is created on exactly one condition, that the plugin returns true from requires_xla_coordinator. When it is, the file reads the rank and world size from PJRT_LOCAL_PROCESS_RANK, LOCAL_RANK, RANK, PJRT_LOCAL_PROCESS_COUNT, LOCAL_WORLD_SIZE and WORLD_SIZE (the torchrun names are the fallbacks), reads the address from MASTER_ADDR and XLA_COORDINATOR_PORT, starts the coordinator, and passes the resulting store into GetCApiClient under the key prefix "pjrt:". None of the plugins shipped in the tree override requires_xla_coordinator, so on a stock install that block is skipped and the coordinator stays null.

How PJRT_DEVICE Becomes a Client: pjrt_registry.cpp, Line by Line

Questions

Related explainers