PyTorch/XLA

PyTorch/XLA (torch_xla) is a Python package that uses the XLA deep learning compiler to connect the PyTorch framework to Cloud TPUs and other XLA devices. The repository is jointly operated and maintained by Google, Meta and individual contributors. The project notes that once TorchTPU is public it will replace PyTorch/XLA.

What It Is

PyTorch/XLA is a compiler backend for PyTorch. Existing PyTorch models and training loops are moved to an XLA device with minimal code changes, such as moving tensors and models to the xla device and wrapping steps in torch_xla.step(). It is installed with pip, for example pip install torch torch_xla[tpu], and also ships as Docker images and nightly wheels.

Training Modes and Distribution

The docs describe single-process training, multi-process training launched with torch_xla.launch, and SPMD, where one Python process controls all TPUs on the system. Distributed support covers Distributed Data Parallel, two FSDP approaches (one process per accelerator, and SPMD-based), distributed checkpointing and Torch Distributed integration through an xla process-group backend.

Advanced Techniques

The documentation covers custom kernels via Pallas on TPU and via Triton on GPU, export to StableHLO, automatic mixed precision, dynamic shapes, TorchDynamo integration and quantized operations. It also documents scan and scan_layers for repeated layers, while_loop for memory use, and @assume_pure to speed up tracing. Troubleshooting guides address tracing versus execution time and sources of recompilation.

Current Status

A notice in the README and docs points to the TorchTPU announcement, stating that TorchTPU will replace PyTorch/XLA once it is public. An October 2025 RFC proposed a more native direction for PyTorch on TPU. The latest release listed is v2.9.0.