AWS's TorchNeuron Runs PyTorch on Trainium With .to('neuron'). Shapes Must Be Static

• Book: RAG Pocket Guide
• The series: Pocket Guides for Developers — 8 short books on agents, RAG, prompts, observability, events, databases and system design — the whole series
• My project: Hermes IDE | GitHub — an IDE for developers who ship with Claude Code and other AI coding tools
• Me: xgabriel.com | GitHub

---

model = MyModel().to('neuron'), with the comment # the only change. That is the headline code sample in the PyTorch blog post where AWS introduced TorchNeuron on October 19, 2026.

TorchNeuron is a native PyTorch backend for AWS Trainium chips. AWS says it goes open source under Apache 2.0 as a public beta on October 20, at PyTorch Conference North America. Amazon's own announcement calls the switch a change with "no rewrite". TechCrunch framed it as Amazon's biggest shot yet at Nvidia's software lead.

The device string really is that short. The fine print in the repository README is where your week goes, so read it before you book Trainium capacity.

What AWS shipped

Until now, running PyTorch on Trainium took "a significant amount of manual work," as SiliconANGLE put it. TorchNeuron registers Trainium through PrivateUse1, PyTorch's official extension point for custom hardware, so torch.nn, torch.autocast, torch.distributed and torch.profiler work on a device named neuron.

Per the PyTorch blog post, it supports:

• eager execution, plus an experimental "adaptive eager" mode that cuts launch cost and fuses operations
• torch.compile with a neuron backend, for production
• DDP, FSDP2, tensor and pipeline parallelism, DTensor and DeviceMesh
• custom kernels through the Neuron Kernel Interface (NKI)
• Trainium2 and Trainium3

The announcement says Trainium is now a supported device upstream in Hugging Face's Transformers, Diffusers, TRL and kernels libraries. Loading a model takes device_map="neuron". The TorchNeuron docs add that the torch.compile backend does not use TorchInductor. It compiles each FX graph into a single Neuron executable.

The one line, and the lines around it

Here is AWS's eager-mode example from the announcement, unchanged:

# Run a training loop on Trainium
model = MyModel().to('neuron')   # the only change
optimizer = torch.optim.AdamW(model.parameters())
for batch in dataloader:
    optimizer.zero_grad()
    loss = criterion(model(batch), targets)
    loss.backward()
    optimizer.step()

For speed you compile. The README quick start passes the backend and turns dynamic shapes off:

compiled_model = torch.compile(model, backend="neuron", dynamic=False)

Getting there is not one line. You need a Trn2 or Trn3 instance and Python 3.12, 3.13 or 3.14, per the README. The manual install guide pulls packages from AWS's own pip index plus PyTorch's CPU wheels. This is its command with one change: I added the torch==2.14.0 pin the README uses.

python -m pip install "torch==2.14.0" \
  "torch-neuronx==2.14.*" "neuronx-cc==2.*" nki \
  --extra-index-url https://pip.repos.neuron.amazonaws.com \
  --extra-index-url https://download.pytorch.org/whl/cpu

The README recommends the Neuron Deep Learning AMI or container instead, where it comes preinstalled. Either way, it tells you to check the install with python -m torch_neuronx.collect_env and look for "device usable from torch" set to yes.

Same package name, two different products

This is the trap worth a sticky note. The package is still called torch-neuronx, and the old one was a different stack. The README says it directly: "Keep the version pins. Earlier torch-neuronx versions are the XLA-based product."

The Neuron docs confirm it: the XLA-based library, versions before 2.14, "remains available as a separate version lineage of the same package." The native backend ships as torch-neuronx for PyTorch 2.14 with Neuron SDK 2.33.

So an unpinned pip install torch-neuronx in an old lockfile, a stale Dockerfile or a tutorial from last year may not give you what this announcement describes. Pin torch==2.14.0 and torch-neuronx==2.14.* together, as in the command above, and treat any Trainium guide written against xla devices as a different product.

What the beta doesn't do yet

The README's known limitations are short and specific, and most of them are about torch.compile:

• torch.compile needs static shapes. Each new input shape compiles again.
• It requires dynamic=False. Only the default mode is supported; other modes run as default with a warning.
• It supports only contiguous tensors; non-contiguous inputs get copied on every call.
• float64 runs as float32, and int64 support depends on the operation.
• TorchAO integration, including MXFP8, is not in this release.

The PyTorch blog post is open about the gaps, too. AWS's August run of PyTorch's upstream test suites produced about 154K outcomes. About 39K were deliberate Neuron skips, and about 22K of those were operators without a native Neuron lowering yet. Features left out of the first release include sparse tensors, NestedTensor, forward-mode autograd, and "dynamic and symbolic shapes."

The Neuron docs also spell out what happens when an op is missing. In eager mode it runs on the CPU with a warning. Under torch.compile, an op the backend can't lower fails the compile by default. A CPU fallback inside a training loop slows every step without failing anything, so run a short eager smoke test and read the warnings before you start a long job.

Static shapes matter most for language models, where every prompt has a different length. The README's advice: "Pad or bucket the inputs to limit the number of shapes." A minimal version of bucketing looks like this:

import torch
import torch.nn.functional as F

BUCKETS = (128, 256, 512, 1024)

def pad_to_bucket(ids: torch.Tensor, pad_id: int) -> torch.Tensor:
    n = ids.shape[-1]
    size = next(b for b in BUCKETS if b >= n)
    return F.pad(ids, (0, size - n), value=pad_id)

Four buckets means at most four compiled sequence lengths. Pad the attention mask the same way, and decide up front what happens to inputs longer than your largest bucket, because next() raises StopIteration there.

Read the numbers as vendor numbers

Every performance claim so far comes from AWS or its customers. Deepgram's CTO, quoted in the PyTorch blog post, says that in the benchmarks Deepgram evaluated, Trainium processes "about 60% more input audio seconds per dollar" than its existing GPU baselines. In TechCrunch's report, Colin Brace, VP of Technology Software at AWS Annapurna Labs, puts Deepgram at "70% more audio per dollar at production-grade accuracy." Neither number comes from an independent benchmark. SiliconANGLE, for its part, attributes its Trainium3 performance figure to Amazon.

TechCrunch also adds context Amazon's posts leave out: PyTorch users can already move apps to AMD GPUs and Apple chips with a configuration change. Brace told the outlet, "The real story is removing the lock-in that kept developers from having a real choice at all."

The useful part for you is narrower than the lock-in debate. If your model already runs on fixed shapes and common ops, trying Trainium now comes down to a pinned install and an afternoon of smoke tests. Dynamic shapes and sparse tensors are another matter: the launch post lists both among the features left out of the first release. I haven't run TorchNeuron myself, so weigh this as a reading of AWS's docs, not a field report.

Have you tried moving a PyTorch workload off CUDA, to Trainium, AMD or Apple silicon? What broke first? Tell me in the comments.

---

If this was useful

I build Hermes IDE, a free, source-available IDE for running Claude Code, Codex and other coding agents side by side, so weigh this with that in mind. If the models you'd move to Trainium are embedding models and rerankers behind a search box, the chip is the smallest decision in that pipeline. The RAG Pocket Guide covers the rest: chunking, retrieval and reranking patterns that decide whether the answers are any good. It's one of 8 short books in the Pocket Guides for Developers series, covering agents, RAG, prompts, observability, events, databases and system design.