Google open-sources TorchTPU, a native PyTorch back end for TPUs
Points and comments are a snapshot, not live.
Google open-sourced TorchTPU, a native PyTorch backend for its TPU hardware.
TorchTPU provides custom ATen kernels and a compilation cache for running PyTorch models on Google TPUs. It supports distributed training (data and tensor parallelism) and includes examples for Llama, ResNet, and UNet. The project integrates with `torch.compile()` for graph-mode execution. The latest release (0.2.0) is available on PyPI for Python 3.11+. Building from source uses Bazel. Documentation and interactive Marimo notebooks are available.
What commenters are saying
Commenters noted the timing with the PyTorch conference and saw TorchTPU as a move to make TPUs more accessible beyond Google Cloud, with some mentioning Google sells TPUs on-prem to large customers. Several pointed out this effectively replaces the older pytorch/xla project. One commenter shared their own open-source libtpu implementation targeting Tenstorrent accelerators. A few questioned whether TPUs are widely available outside Google's ecosystem, with others noting the Pixel phone NPUs are also called TPUs.