Google and Meta release a native PyTorch backend for TPUs

Google and Meta’s PyTorch team introduced TorchTPU, an open-source backend that lets PyTorch programs run on Google’s tensor processing units (TPUs) as a native device, with eager execution and torch.compile. The project began at Google in mid-2025 in collaboration with Meta’s PyTorch team. On the same day Amazon Web Services opened a public beta of TorchNeuron, an equivalent open-source backend for its Trainium chips.

The post says TPUs have historically been associated with the TensorFlow and JAX ecosystems. Google’s earlier route for PyTorch, PyTorch/XLA, which dates from 2020, is described in the post as needing intrusive code changes, suffering long compile times and lacking eager mode. TorchTPU is built, the authors say, with no JAX dependency, and the stated aim is that porting a workload can be as simple as changing model.to("cuda") to model.to("tpu"). The post adds that peak TPU performance will not come from that change alone, because a model’s shapes and sharding may be tuned to GPUs.

The post says TorchTPU is in production for inference and training at some of Google Cloud’s largest customers. In Google’s own compatibility testing, 4,821 of 4,917 of the most-downloaded Hugging Face models (98%) ran without crashing, and 638 of 657 representative architectures (97%) matched GPU numerics. On a Qwen3 235B mixture-of-experts model across 64 TPU v7 chips, PyTorch via TorchTitan reached 1,828 tokens per second against 1,937 for JAX with MaxText. These are Google’s results; the post does not give an independent comparison with GPUs.

AWS’s release gives developers .to('neuron') with eager mode and torch.compile on Trainium2 and Trainium3 instances. Hugging Face, which published its own post a day earlier, said 89% of the 100 most-used text-generation model architectures now run on Trainium out of the box, against 46% with its previous Optimum Neuron library. That figure is Hugging Face’s.

Neither announcement is framed as a challenge to Nvidia; the TorchTPU post does not mention it. Both are aimed at the same practical point, that PyTorch code written for one device string can be pointed at another accelerator with small changes. Whether that alters where models are trained will depend on performance and availability on production workloads, which these posts report only in part.