TorchTPU brings a native PyTorch backend to Google TPUs

The TorchTPU team’s announcement introduces the project on October 20, 2026, while noting that development continues. On a TPU machine with its package installed, importing PyTorch loads the backend automatically. The team illustrates migration by changing the device target from cuda to tpu. It also warns that this change alone cannot deliver peak performance: model shapes, operations and distribution across devices may still reflect GPU assumptions.

The project repository documents installation in a Python virtual environment, with Python 3.11 or later supported. Its release installation command selects torch_tpu 0.2.0 and supplies PyTorch’s CPU package index. Omitting that index installs the CUDA version of PyTorch unless the CPU build is installed separately. For exploration, the repository includes distributed-training and model-inference examples. Its interactive Marimo notebooks run from a TPU virtual machine and open in a browser; the user guide provides their static version.

The public API documentation distinguishes interfaces supported long term from experimental ones. Incompatible changes to stable interfaces follow a deprecation process. Experimental interfaces can change or disappear without that process and emit a one-time warning when used. For example, torch.accelerator.is_available is listed as stable, while torch.tpu.Event is experimental. The page warns that using other TorchTPU APIs is unsupported; PyTorch’s own documentation covers the PyTorch interfaces implemented for TPUs.