# PyTorch/XLA

> Python package that implements the XLA compiler as a backend for PyTorch, enabling training on Cloud TPUs and other XLA devices.

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.

## Features
- XLA backend for PyTorch on Cloud TPUs
- Single-process, multi-process and SPMD training
- Distributed Data Parallel and FSDP support
- Custom kernels via Pallas and Triton
- Torch export to StableHLO
- Automatic mixed precision
- TorchDynamo integration
- Dynamic shapes support
- Quantized operations
- scan, scan_layers and while_loop optimizations
- Distributed checkpointing
- Custom hardware plugins
- Docker images and nightly wheels

## Integrations
PyTorch, Google Cloud TPU, XLA, OpenXLA, Kaggle, Pallas, Triton, StableHLO, TorchDynamo

## Platforms
WINDOWS, MACOS, LINUX, API, DEVELOPER_SDK

## Pricing
Open Source

## Version
v2.9.0

## Links
- Website: https://pytorch.org/xla
- Documentation: https://docs.pytorch.org/xla/release/r2.8/index.html
- Repository: https://github.com/pytorch/xla
- EveryDev.ai: https://www.everydev.ai/tools/pytorch-xla