AWS and Google ship PyTorch backends for Trainium and TPU

AWS released TorchNeuron for Trainium chips, and Google with Meta released TorchTPU for TPUs. Both are open source and plug into PyTorch through the same PrivateUse1 registration hook. Google tested 4,917 Hugging Face models: 98% ran without crashing and 97% of those checked matched GPU numbers.

• TorchTPU supports around 450 PyTorch operations
• Skinny Collectives option costs 30–40% end to end
• 64-chip test: 1,828 vs 1,937 tokens/s against JAX/MaxText
• TorchTPU drops the old JAX path for a native backend