PJRT Runtime (Beta)
PyTorch/XLA is migrating from XRT to the new PJRT runtime. PJRT is a better-maintained stack, with demonstrated performance advantages, including, on average, a 35% performance for training on TorchBench 2.0 models. It also supports a richer set of features enabling technologies like SPMD. In the PyTorch/XLA 2.0 release, PJRT is the default runtime for TPU and CPU; GPU support is in experimental state. The PJRT features included in the PyTorch/XLA 2.0 release are:

TPU runtime implementation in libtpu using the PJRT Plugin API improves performance by up to 30%
torch.distributed support for TPU v2 and v3, including pjrt:// init_method (Experimental)
Single-host GPU support. Multi-host support coming soon. (Experimental)
