Google open-sources TorchTPU, a native PyTorch backend for TPU accelerators

Google open-sources TorchTPU for TPU hardware

Google has released TorchTPU, a PyTorch backend for its Tensor Processing Units, as an open-source project under the google-pytorch organization on GitHub. The release, which surfaced on the Hacker News front page on 20 October, is already installable from PyPI at version 0.2.0 and carries an Apache 2.0 license.

According to the project's README, TorchTPU supplies custom ATen kernels — the operator implementations sitting underneath PyTorch's Python API — so that PyTorch models and individual operations can execute on TPU devices. A compilation cache retains optimized forms of the XLA computations the backend generates, and the project integrates with torch.compile() for graph-mode execution.

What ships in the box

Beyond the kernels, the repository lists support for distributed training using both data parallelism and tensor parallelism, together with worked examples covering Llama, Qwen, ResNet, DLRM, minGPT and UNet. Debugging and benchmarking utilities are included as well.

The documentation doubles as a set of interactive notebooks. A user guide is hosted on Read the Docs as a static rendering of Marimo notebooks, a Jupyter-like tool, and the README notes these can be edited and run in a browser with marimo edit. Running them requires a TPU VM, with a v6e Google Compute Engine instance flagged as the simplest option.

Installation and building

The project supports Python 3.11 and newer, mirroring PyTorch's release compatibility matrix. Installation goes through pip, but the README warns that you must pass PyTorch's CPU wheel index as an extra index URL; otherwise pip pulls in the default CUDA build of torch, and the project directs users to the CPU build instead.

Building from source goes through Bazel, installed via the Bazelisk wrapper, and the README documents several sharp edges. Release wheels are always built in optimized mode, builds fail outright without a required wheel configuration flag, and developers without Remote Build Execution credentials need to append a further flag to strip those settings back out. Contributors are welcome, but must sign Google's Contributor License Agreement and go through code review before changes land.

Why it matters

PyTorch is the default working environment for a large share of machine-learning researchers and engineers, while TPUs have historically been most convenient to reach from Google's own framework stack. A purpose-built, open-source backend that plugs into PyTorch at the ATen layer — and works with torch.compile() — lowers the friction of moving existing PyTorch training and inference workloads onto Google's accelerators instead of rewriting them.

The Apache 2.0 licensing, public examples for current-generation models such as Llama and Qwen, and an explicit contribution process also suggest Google wants outside developers building against this path rather than treating it as an internal tool. At version 0.2.0 the project is still early, but it arrives at a moment when competition among accelerator ecosystems is intensifying, and framework accessibility is one of the levers that decides where workloads actually run.