Google open-sources TorchTPU, PyTorch's native backend for TPUs

What TorchTPU is

Google has open-sourced TorchTPU, a PyTorch backend for its TPUs (Tensor Processing Units), in the google-pytorch/torch_tpu repository on GitHub, under the Apache 2.0 license. The project's official description sums up the idea:

This project provides a PyTorch backend for Google's Tensor Processing Units (TPUs), enabling users to run PyTorch models and operations on TPU devices.

Project description, google-pytorch/torch_tpu repository

In practice, the project implements ATen kernels, the operator layer that PyTorch uses to dispatch tensor operations, specifically targeted at TPU. It's the same layer that the CPU and CUDA backends already use, which brings the experience of running models on TPU closer to that of someone who already trains on GPU day to day.

The path that already existed: pytorch/xla

Before TorchTPU, the established way to run PyTorch on TPU was the pytorch/xla library, which exposes an xla device and works with lazy tensors: operations are only actually executed when the code explicitly calls xm.mark_step(). This model works, but it requires adapting the training loop and understanding an execution semantics different from the eager mode that developers already know from CPU and GPU.

TorchTPU attacks this friction point at the root: by implementing the kernels directly at the ATen layer, the goal is for common PyTorch code to run on TPU without rewriting training around lazy tensors. The README lists, among its features, integration with torch.compile() for graph mode execution, the same compilation mechanism that already exists in PyTorch 2.x for CPU and GPU.

XLA hasn't disappeared, it just moved

An important correction is due here regarding the expectation that native TPU support means goodbye to XLA: it's not quite like that. The README itself lists among its features a "compilation cache for optimized execution of XLA computations," and the build instructions mention that the --config=wheel_common flag "routes the shared XLA/MLIR backend" so that it gets compiled into a single libxla_base.so.

In short: TorchTPU changes the interface the developer uses (native ATen kernels, without explicit lazy tensors), but the compiler that actually generates code for TPU hardware is still the XLA/MLIR stack underneath. It's a change in the abstraction layer, not the elimination of XLA from the stack.

How to install and what's already ready

The latest public version is 0.2.0, available on PyPI:

pip install torch_tpu==0.2.0 --extra-index-url https://download.pytorch.org/whl/cpu

The project requires Python 3.11 or later, following the same compatibility policy as PyTorch. Among the features already documented in the README:

• Support for distributed training, with Data Parallelism and Tensor Parallelism
• Ready-made examples for Llama, Qwen, ResNet, DLRM, minGPT, and UNet
• Compilation cache to speed up repeated runs
• Debugging and benchmarking utilities
• Interactive user guide in Marimo notebooks, a tool similar to Jupyter

The notebooks run with marimo edit file.py, but they need to be executed on an actual TPU VM. The README itself recommends a v6e instance on GCE as the simplest path for anyone who wants to test it, which means you can't just clone the repository and run it locally: you need to provision TPU hardware on Google Cloud.

Building it yourself requires Bazel and patience

Anyone who wants to compile TorchTPU from source, instead of installing the ready-made wheel, runs into a Bazel-based build system that's much more elaborate than pip install suggests. The basic command is:

bazel build --config=wheel_common //ci/wheel:torch_tpu_wheel

The --config=wheel_common flag is not optional, the build fails without it, because it's the one that enables the --//:wheel_build=True setting needed to avoid duplicate registrations when import torch runs. Anyone without Remote Build Execution (RBE) credentials needs to add --config=no_rbe after wheel_common, in that order, because the second config overrides the remote cache flags from the first.

The repository has 4,466 commits recorded at the time of publication, a volume that suggests internal development at Google for a good while before the code was open-sourced. Even so, parts of the documentation are incomplete: the Python dependencies section of the README itself just has a "TBD" in place of the list.

What remains open

The README doesn't provide any performance numbers, nor a cost or speed comparison against pytorch/xla, against GPU, or between TPU versions. For anyone deciding between training a large model on TPU or GPU on GCP, the concrete change TorchTPU delivers so far is one of ergonomics and engineering effort, not a publicly proven FLOPs count or dollars-per-hour figure.

With the project still at version 0.2.0 and dependency documentation marked as pending, it makes sense to treat TorchTPU as an early-stage bet: it's worth following subsequent releases and testing it on a small project before migrating an entire training pipeline to it. The repository is at https://github.com/google-pytorch/torch_tpu and the user guide is at https://torchtpu.readthedocs.io/.