PyTorch Gets Native TPU and Trainium Backends in the Same Week -- Google's TorchTPU and AWS's TorchNeuron Replace the Lazy-Tensor XLA Path, and Hugging Face Retires Its Workaround Libraries

For years, running PyTorch on Google TPUs or AWS Trainium meant going through PyTorch/XLA and its lazy tensors. Two releases on October 19-20, 2026 replace that path with native backends.

TorchTPU

Google's google-pytorch/torch_tpu repository, Apache-2.0, published version 0.2.0 to PyPI on October 19, and Hacker News picked it up the next day (269 points). The README lists "custom ATen kernels", a compilation cache for XLA computations, data and tensor parallelism, and torch.compile() integration, with examples from Llama and Qwen to DLRM and UNet. Hugging Face's October 20 post explains what changes for users: "With TorchTPU, a native PyTorch backend, TPU is now just another device." Under the old route, "Tensors were lazy (the tensors, not you), so your code didn't always run when you thought it did," and "Small changes could trigger graph recompilations." Now Transformers takes device_map="tpu", Diffusers pipelines take pipe.to("tpu"), and Optimum TPU "is not required anymore". Kernels written in Pallas can be loaded from the Hub. With a ragged-dot MoE kernel from Google's Tokamax library on a TPU v6e-8, "At batch size 64, decoding is 21.5× faster on Qwen3-30B-A3B and 12.8× faster on Gemma 4 26B-A4B" than the default implementation. Hugging Face is candid about maturity: "Of the tests that run on TPU, 96.5% pass." It has filed 34 issues, 19 of them already fixed.

TorchNeuron

Hugging Face's October 19 post says TorchNeuron would enter public beta on October 20 via pip install torch-neuronx. It exposes neuron as a native device with eager mode, torch.compile, DTensor and FSDP2. The coverage figure: "89% of the top 100 most-used text-generation model architectures are supported out of the box, compared to 46% with Optimum Neuron v0.4.5". Hugging Face now recommends using Transformers, Accelerate, PEFT and TRL directly rather than Optimum Neuron. Some constraints remain. Static shapes are still advised; a standard TRL fine-tuning script turns off dynamic shapes, and TRL 1.16 or later pads Neuron batches to multiples of 128 tokens to limit recompilation.

Commenters made the obvious connection. One pointed to AWS's parallel work; another said they were building an open libtpu for Tenstorrent chips that "can already run most of torchtpu", a claim we have not checked.

What remains uncertain. Neither post compares cost or speed against GPUs, against the old XLA path, or between the two backends. The 21.5x figure measures an optimised kernel against an unoptimised default, not against a competitor. TorchTPU is at version 0.2, and TorchNeuron is in beta.

Analysis: CUDA's advantage at the framework level is shrinking. The remaining differences sit in kernels (Pallas, NKI and CUDA or Triton), in compilers, and in price. Teams can prepare now at little cost: keep device selection in configuration rather than code, avoid CUDA-only calls in model code, and bucket sequence lengths so compiled graphs get reused. Keep custom kernels behind a swappable layer such as the Kernels Hub. Then benchmark cost per token on your own workload before moving anything. That last step matters, because neither launch publishes a number that answers the question buyers care about.