PyTorch/XLA documentation

PyTorch/XLA is a Python package that uses the XLA deep learning compiler to connect the PyTorch deep learning framework and Cloud TPUs.

Learn about Pytorch/XLA

• Pytorch/XLA overview
• PyTorch on XLA Devices
• PyTorch/XLA API
• Dynamic shape
• Eager Mode + Compile API
• PJRT Runtime
• Troubleshoot

Learn about accelerators

• Learn about TPUs
• Learn about GPUs

PyTorch/XLA features

• Custom Kernels via Pallas
• Torch Export to StableHLO
• Custom GPU Kernels via Triton
• Guide for using scan and scan_layers

Improve Pytorch/XLA workload performance

• Automatic Mixed Precision
• PyTorch/XLA SPMD User Guide
• PyTorch/XLA SPMD advanced topics
• Distributed Checkpointing
• Running SPMD on GPU
• How to do DistributedDataParallel(DDP)
• TorchDynamo integration in PyTorch XLA
• Optimize memory utilization using while_loop
• Fully Sharded Data Parallel in PyTorch XLA
• Fully Sharded Data Parallel using SPMD
• Quantized Operations for XLA (Experimental feature)
• Source of recompilations in Pytorch/XLA

Contribute to Pytorch/XLA

• Configure a development environment
• Codegen migration Guide
• OP Lowering Guide
• Custom Hardware Plugins
• Bazel in Pytorch/XLA