Google's TorchTPU is on GitHub: a PyTorch backend that asks you to change one word, the device name, to run on a TPU
Version 0.2.0 ships custom ATen kernels, three eager modes and an XLA-based torch.compile path; Google claims Fused Eager runs 50% to 100+% faster than Strict Eager, but publishes no GPU comparison.
it should feel like PyTorch
Comprehensive take
Axion take: a credible attempt to make TPUs a drop-in PyTorch device, backed so far by design intent and a 0.2.0 tag. Try it; do not plan a migration on it until benchmarks arrive.
Use cases
Porting existing PyTorch training to TPUs
Google's post says scripts should migrate with minimal changes, with device initialisation set to "tpu".
Debugging before scaling
Debug Eager syncs after every op to surface shape mismatches, NaNs and out-of-memory errors; Fused Eager is the fast path.
Who should wait
Teams needing published benchmarks, documented TPU generation support, or dynamic shapes: the README leaves dependencies as TBD and the roadmap still lists these.
Specs
| Release | torch_tpu 0.2.0 on PyPI source |
|---|---|
| License | Apache 2.0 source |
| Requirement | Python 3.11+ source |
| Fused Eager vs Strict Eager | 50% to 100+% faster (Google claim) source |
| Compiler path | Torch Dynamo graph to XLA via StableHLO source |
| Distributed support | DDP, FSDPv2, DTensor source |
Insights
- Google's own only numeric performance claim is Fused Eager at 50% to 100+% over Strict Eager.
- Unlike PyTorch/XLA, which handled pure SPMD code, TorchTPU is designed for divergent MPMD execution.
- The README lists dependencies as TBD and does not name supported TPU generations.
Google's TorchTPU is on GitHub, and its pitch is that moving a PyTorch script to a TPU should be a one-word change. The repository at google-pytorch/torch_tpu lists a 0.2.0 release on PyPI, an Apache 2.0 license, custom ATen kernels, a compilation cache, data and tensor parallelism, and example models including Llama, Qwen, ResNet, DLRM, minGPT and UNet. Google's design post from 7 April 2026 states the goal as "it should feel like PyTorch", with device initialisation changed to "tpu".
How it plugs in
According to Google's developer blog, TorchTPU uses PyTorch's PrivateUse1 interface, with no subclasses or wrappers, so developers keep working with ordinary PyTorch tensors. It offers three eager modes. Debug Eager runs one operation at a time and syncs with the CPU after each, which the post says is slow but useful for catching shape mismatches, NaNs and out-of-memory errors. Strict Eager dispatches single operations asynchronously like default PyTorch. Fused Eager automatically fuses operations into larger chunks and, Google says, is 50% to 100+% faster than Strict Eager with no user setup.
For full-graph work, TorchTPU hooks into torch.compile, captures the graph with Torch Dynamo, and uses XLA rather than Torch Inductor as the compiler, mapping PyTorch operators to StableHLO. Distributed training covers DDP, FSDPv2 and DTensor, and the post says it is designed to support divergent (MPMD) execution, unlike its predecessor PyTorch/XLA, which handled only pure SPMD code. Pallas and JAX kernels can be injected with a decorator.
What is still missing
The README is thin. Its dependencies section reads "TBD", it does not list supported TPU generations, and it carries no performance numbers beyond a warning that unoptimised builds can degrade speed. The only hardware hint is that example notebooks must run on a TPU VM, with a v6e instance called the easiest. The blog's own numeric claim is the Fused Eager gain; it gives no end-to-end training or inference benchmarks against GPUs. Its 2026 roadmap still lists dynamic shape support, precompiled TPU kernels, deeper vLLM and TorchTitan integration, and validated linear scaling to full Pod size.
Who should try it
Teams with existing PyTorch code and access to TPU capacity on Google Cloud have the most to gain: the install is a pip command from the PyTorch CPU wheel index, and Python 3.11 or newer is required. Teams that need proven large-scale results should wait for benchmarks and the roadmap items above.
Axion take
The strategic bet is clear: lower the cost of leaving CUDA by making TPUs behave like a normal PyTorch device. The evidence so far is design intent and a 0.2.0 tag, not measured results. Treat it as worth a weekend experiment, not yet a migration plan.