TorchTPU vs TorchNeuron: running PyTorch on Google TPUs and AWS Trainium compared (October 2026)

TorchTPU vs TorchNeuron is the new choice for PyTorch developers who want to run on cloud AI chips that aren't Nvidia GPUs. On October 20, 2026, Google open-sourced TorchTPU, a native PyTorch backend for its Tensor Processing Units, a day after Amazon shipped TorchNeuron, its native backend for Trainium, as a public beta. Both let ordinary PyTorch code target the custom accelerators in Google's and Amazon's AI datacenters by changing the device name, and both replace an older route through PyTorch/XLA. This comparison sets out, as of October 23, 2026, how they differ.

This comparison is based on Google's and the PyTorch Foundation's October 20 announcements, the TorchTPU documentation (version 0.2.0), repository and PyPI page, AWS's Neuron 2.33.0 release notes and the torch-neuronx repository, checked October 23, 2026. It covers what the docs say, not how fast either runs. For the TorchNeuron launch itself, see our TorchNeuron story.

The short version

• Same idea, different chips. Both register a new PyTorch device through PrivateUse1: tpu for Google TPUs, neuron for AWS Trainium2 and Trainium3.
• Status: TorchTPU is version 0.2.0. Google's docs and announcements give it no alpha, beta or GA label, though its PyPI package metadata says "Development Status :: 4 - Beta"; Google's FAQ calls it "the new official standard" for PyTorch on TPUs. TorchNeuron is a public beta in AWS's Neuron 2.33.0.
• License: Apache 2.0 for both.
• PyTorch versions: TorchNeuron names 2.13 and 2.14. TorchTPU's docs give no supported range; its 0.2.0 package on PyPI requires PyTorch 2.11.0 to 2.15.0.
• Compiled mode: torch.compile(model, backend="tpu") or backend="neuron"; both default to static shapes.
• Distributed: both ship their own c10d process-group backend (tpu_dist, neuron) and support DTensor and DDP; FSDP2 is listed as experimental, with low performance, for TPUs and as supported for Trainium.
• The old path: Google Cloud says "TorchTPU replaces PyTorch/XLA." AWS ended PyTorch/XLA training support from Neuron 2.30.0 but keeps it for inference.

How do TorchTPU and TorchNeuron compare?

Last verified: October 23, 2026 (every cell with its source link is in our data file D31-torchtpu-vs-torchneuron.csv)

Two cells need context. TorchTPU's GitHub page reports its license as "Other" because the license file adds an MIT section for one example folder (examples/mingpt/impl); the README and pyproject.toml declare Apache 2.0. And the PyTorch range and the "Beta" label both come from the 0.2.0 package's metadata on PyPI (torch>=2.11.0,<=2.15.0 and "Development Status :: 4 - Beta"), which matches the repository's v0.2.0 tag. Google's docs and announcements for 0.2.0 give neither; the docs' API reference says only that its torch.accelerator APIs need PyTorch 2.11 or later. The main branch, already numbered 0.2.1, drops the upper bound and requires PyTorch 2.11.0 or later.

What code do you change?

Mostly the device and the distributed backend. TorchTPU's introduction says porting from GPUs "starts with changing device="cuda" to device="tpu" and backend="nccl" to backend="tpu_dist"." The PyTorch blog post on TorchTPU gives this before-and-after:

# GPU
device = torch.device("cuda")
dist.init_process_group(backend="nccl")

# TPU
device = torch.device("tpu")
dist.init_process_group(backend="tpu_dist")

For TorchNeuron, the PyTorch blog's example is model = MyModel().to('neuron'), and the release notes use init_process_group(backend="neuron").

Both teams also point to PyTorch's device-agnostic API. TorchTPU's docs say the backend works with torch.accelerator.current_accelerator(), and the TorchNeuron blog says code written against torch.accelerator "runs unmodified." Hugging Face integration differs slightly: Google's post shows device_map="tpu" in Transformers v5, while the TorchNeuron blog says device_map="auto" detects Neuron.

The harder change is from the older XLA path. AWS calls moving off torch_xla and xm.xla_device() a breaking change. The PyTorch blog's TorchTPU post says PyTorch/XLA "required intrusive source-code modifications to explicitly define compilable regions," which TorchTPU is designed to avoid.

Neither vendor promises full speed from the device swap alone. The TorchTPU post says "one cannot achieve peak TPU performance with that simple change," but the model "should run out of the box."

How do eager mode and torch.compile differ?

Both run eager by default, one operation at a time, and both offer torch.compile for production.

• TorchTPU has three eager modes. Strict, the default, dispatches one ATen operator at a time asynchronously. Debug (TORCH_TPU_LAUNCH_BLOCKING=1) runs synchronously. DeferAndFuse (TORCH_TPU_DEFER_AND_FUSE=1) batches operations, which the PyTorch blog says "can be 50% or more faster than Strict Eager." That figure is vendor-reported with no stated method. Its compiled mode hands the whole FX graph to Google's XLA compiler as one StableHLO program and doesn't use TorchInductor.
• TorchNeuron runs eager by default, with an experimental "adaptive eager" mode switched off. Its torch.compile backend requires dynamic=False and mode="default", and AWS warns that "each new input shape can require another compilation."

Static shapes are the shared limit. Google's docs say torch.compile sets dynamic=False by default on TPUs, and the PyTorch blog post explains that bounded dynamism, which dynamic shapes need, is "not enabled by default yet."

Which distributed training APIs work?

Last verified: October 23, 2026

TorchTPU's docs say standard torch.distributed calls work in both eager and compiled modes, launched with torchrun after setting topology variables. By default each collective runs as its own small program, which the PyTorch blog says can cost "e.g., 30-40% end-to-end" until you fuse them with its @spmd_safe decorator.

What are the known limits?

TorchTPU, per Google's docs and the PyTorch blog:

• about 450 ATen operators implemented;
• static shapes by default, with recompilation when shapes change;
• a default medium matmul precision (BF16 multiply, FP32 accumulate), unlike CUDA's highest; the blog says high or highest cost about 3x or 6x the matrix-unit compute;
• no native FlexAttention yet;
• custom Pallas kernels still need JAX installed to author, though the runtime has no JAX dependency;
• the README's "Dependencies" section still reads "TBD."

TorchNeuron, per AWS's release notes and README:

• no torch.sparse, torch.func, torch.nested or torch.distributions;
• no TorchAO integration or composite MXFP8 support;
• dynamic shapes "still in development";
• only contiguous tensors, and float64 runs as float32 (README);
• Ubuntu 22.04 and 24.04 wheels "Planned for the next release."

How mature are the docs and ecosystems?

TorchTPU arrives with a dedicated documentation site for version 0.2.0: getting-started guides for TPU VMs and GKE, a feature-set page that separates "Validated" from "Experimental," an API reference, an FAQ, and case studies including Meta's open-weight Muse Glimmer 30B. Google also open-sourced vLLM-TorchTPU, SGLang-TorchTPU and a TorchTPU integration of TorchTitan the same day, and lists Ray, TorchRec and Hugging Face integrations.

TorchNeuron's documentation lives inside AWS's broader Neuron docs, with release notes and a "What is TorchNeuron?" overview. As our TorchNeuron story noted, AWS pages gave different minimum PyTorch versions as the backend moved from private to public beta. AWS lists TorchTitan, Hugging Face Transformers, Diffusers and trl among supported libraries.

Both vendors publish performance figures, all vendor-reported and not comparable with each other. Google says TorchTitan reached up to 52% model FLOPs utilization on Qwen 3 8B across up to 8,192 Ironwood TPUs, giving no method. The TorchNeuron blog post quotes Deepgram's claim of about 60% more audio processed per dollar than its GPU baselines. Neither was run under shared conditions, so they say nothing about which chip is faster.

Recommendation by use case

Last verified: October 23, 2026

The bigger context is cloud chip competition. TechCrunch's October 20 headline, by Julie Bort, put it as "Google has joined Amazon in attacking Nvidia's stronghold with AI app developers." Software has long been the obstacle to using custom chips, and both backends aim at it. The capacity behind them is tracked in our gigawatt datacenter analysis, and Google's other bid for enterprise AI buyers that week, selling OpenAI's models, is covered in our OpenAI on Google Cloud story.

What's confirmed, reported and uncertain

• Confirmed (official docs, repos and posts): every cell in the tables above, as of October 23, 2026.
• Conflicting: TorchTPU's GitHub page shows the license as "Other," while the README and package metadata say Apache 2.0; the license file adds MIT terms for one example folder. Google's FAQ says "PyTorch/XLA is deprecated," while Google Cloud's PyTorch page says TorchTPU "replaces" it and doesn't use the word. TorchTPU's docs and announcements give no status label, while its PyPI metadata says "Beta."
• Vendor-reported: all performance numbers, including DeferAndFuse speedups, MFU figures and Deepgram's cost claim.
• Uncertain: whether Google supports the whole PyTorch 2.11.0 to 2.15.0 range that TorchTPU 0.2.0's package accepts, when TorchNeuron leaves beta, and when either supports dynamic shapes without recompilation.