On October 20, 2026, Google released "TorchTPU" as open source. It is an execution backend that lets the AI development framework PyTorch run on Google's own AI chip, the TPU. Announced at PyTorch Conference 2026, it allows training and inference to run on TPUs without major rewrites of existing PyTorch code. Around the same time, Amazon Web Services (AWS) began offering a public beta of "TorchNeuron," which it developed for its own AI chip, Trainium. What the two efforts have in common is that they let developers keep the AI models and software they already know while making it easier to choose hardware other than NVIDIA GPUs. However, being able to run existing code is not the same as achieving sufficient performance in real services, so developers need to understand the constraints around compilation and supported operations.
Google releases TorchTPU, running existing PyTorch code on TPUs
TorchTPU is a native backend for TPUs, developed by Google together with Meta's PyTorch development team. In PyTorch, developers describe a model's structure and computations, and a backend executes the operations on the actual hardware. In its official blog post of October 20, Google explained that developers can move existing workloads to TPUs with only minor code changes while continuing to use the familiar PyTorch APIs. The aim is to spare developers from having to learn TPU-specific programming techniques from scratch.
TorchTPU supports both an eager mode, suited to development and debugging, and compiled execution, which prioritizes performance. Eager execution runs individual operations in order so that developers can check how a model behaves. In production, torch.compile can be used to optimize computations as a group and make efficient use of the TPU's compute capacity. The design preserves flexibility during development while aiming for higher throughput in operation.
Integration with major AI development libraries is also progressing. In Hugging Face's Transformers v5, TPUs can now be used by specifying device_map="tpu" as the model's placement target. Diffusers, which is used for image generation and other tasks, also supports eager execution, torch.compile, and tensor parallelism on TPUs. Google has additionally built a system that automatically runs Transformers tests on TPUs every night. This is not merely about adding support for new hardware; it is an effort to keep verifying compatibility after Transformers is updated.
TorchTPU support is also spreading to the software that handles AI model inference. Google worked with Inferact to develop a TPU backend for vLLM and with RadixArk to build a backend for SGLang. Both keep the existing worker configuration, scheduler, and API server, while optimizing the hardware-level computation for TPUs. This makes it easier to adopt TPUs without completely rebuilding the systems used to serve AI models to many users.
In training, TorchTPU works with TorchTitan, which is designed for large-scale distributed training, and supports training models across multiple TPUs. TorchTPU has also been integrated into the distributed processing framework Ray, making it possible to manage workflows that combine CPUs, GPUs, and TPUs in one place. Google also released a "TPU agent suite" on GitHub as an early preview, which helps with migrating to TPUs, diagnosing behavior, and tuning inference performance. However, it is at a different development stage from a formal general release, so its available features and maturity need to be evaluated separately.
AWS releases TorchNeuron, treating Trainium as a standard PyTorch device
AWS is likewise building an environment in which its own AI chips can be used directly from PyTorch. According to an announcement on the official PyTorch blog dated October 19, 2026, AWS began offering "TorchNeuron" for Trainium2 and Trainium3 as a public beta from the following day, the 20th. TorchNeuron is an open-source backend released under the Apache License 2.0, and it lets Trainium be treated as a standard PyTorch device. The goal is to run AI models on Trainium without major changes to existing PyTorch code.
TorchNeuron is built on "PrivateUse1," an extension mechanism PyTorch provides for custom hardware. Using it, TorchNeuron registers a device called neuron, making Trainium available through PyTorch's standard APIs. For supported models, placing the model and input data with .to("neuron") enables training and inference in eager mode. To improve performance through compilation, developers either specify the Neuron backend, as in torch.compile(model, backend="neuron"), or register it in advance as the default backend.
There were already ways to run PyTorch models on Trainium, but they came with particular constraints. According to Hugging Face's explanation, the earlier Optimum Neuron was built on PyTorch XLA, which assumed lazy execution and fixed input shapes, so developers had to be mindful of how the compiler behaved. To absorb that complexity, Optimum Neuron maintained its own model-specific implementations and training code. With TorchNeuron, Trainium is treated as a standard PyTorch device, so libraries such as Transformers, Diffusers, Accelerate, PEFT, and TRL can be used directly, and Hugging Face now recommends using each library without going through Optimum Neuron.
The significance of this change goes beyond simpler device specification. If the need to maintain Trainium-specific implementations for each AI model shrinks, it becomes easier to try other hardware while picking up updates to the libraries developers already use. AWS has also proposed changes to make CUDA-dependent code, such as that remaining in PyTorch's own tests, independent of any particular hardware. That said, TorchNeuron does not automatically port all CUDA-specific code to Trainium. It is an effort to increase the portion that relies on PyTorch's standard APIs and reduce hardware-specific implementations.
The difference between Google and AWS lies in production constraints
TorchTPU and TorchNeuron are both technologies that make it easier to use proprietary AI chips through PyTorch's standard APIs. However, they differ in supported hardware, execution methods, and the constraints disclosed so far. Based on both companies' announcements and the AWS Neuron release notes, the main characteristics compare as follows.
| Comparison item | Google "TorchTPU" | AWS "TorchNeuron" |
|---|---|---|
| Release status and target | Released as open source for TPUs | Public beta for Trainium2 and Trainium3 |
| Basic usage | Uses existing PyTorch APIs. In Transformers v5, device_map="tpu" can be specified |
Place the model and inputs on the neuron device and use standard PyTorch APIs |
| Development and debugging | Supports eager mode | Supports eager mode |
| Performance optimization | Supports torch.compile. Custom kernels via Pallas also available |
Supports torch.compile. Custom kernels via NKI can also be written |
| Constraints and conditions disclosed so far | TPU agent suite is an early preview | Compilation in the first version assumes fixed input shapes. In eager mode, unsupported operations can fall back to the CPU |
This comparison organizes the features and availability conditions the two companies announced in October 2026; it is not a comparison of performance or price using the same model. Nor do the constraints explicitly stated for AWS's TorchNeuron necessarily apply to Google's TorchTPU. What the two share is a design philosophy of keeping the application-facing APIs as unchanged as possible and leaving hardware-specific processing to the backend and dedicated kernels. This lowers the barrier to trying other AI chips, but it does not eliminate the tuning needed to get sufficient performance in production.
Model coverage rises to 89%, but performance evaluation calls for caution
As a concrete figure showing the improvement from TorchNeuron, Hugging Face points to expanded model coverage. According to the company, among the 100 most frequently used text-generation model architectures, the share that can be used without additional dedicated implementations rose from 46% with Optimum Neuron v0.4.5 to 89% with TorchNeuron. This means a broader range of models than before can now be tried on Trainium. However, the figure applies only to a specific set of 100 architectures and does not mean that 89% of all AI models work. Support should also be distinguished from running at high speed.
A constraint that particularly needs attention concerns input shapes at compile time. According to the TorchNeuron release notes from AWS, torch.compile in the first version requires specifying dynamic=False, and dynamic input shapes are not yet natively supported. As a result, a change in the shape of the input data can trigger a new compilation. AWS recommends reducing the number of distinct shapes to be compiled, for example by padding inputs to a fixed length or grouping inputs of similar size.
For example, in an AI service that processes large volumes of text of varying lengths, input shapes tend to change from one text to the next. If compilation occurs each time, users may wait longer for a response even when inference itself is fast. Hugging Face's published inference examples also recommend a setup using a fixed-size cache to avoid frequent recompilation. In its training examples, automatic dynamic handling of input shapes by PyTorch Dynamo is disabled, and asynchronous weight loading, which is not currently supported, is turned off. Even when standard model definitions and training code can be used, runtime settings must be checked to run stably on Trainium.
How unsupported operations are handled is another point that cannot be overlooked in performance evaluation. TorchNeuron's eager mode includes a mechanism that switches operations that cannot run on Trainium over to the CPU. In some cases this allows a model to run to completion, but when the switch to the CPU occurs, data transfers and waiting may affect performance. AWS says the operations that fell back can be checked with torch_neuronx.get_fallback_ops(). Even if a model runs correctly, not all operations necessarily ran on Trainium, so it is not appropriate to treat speeds measured without checking where operations executed as the AI chip's own performance.
It is also necessary to measure separately the compilation time on the first run and the run time on later runs, when compilation results can be reused from the cache. Google offers Pallas and AWS offers the Neuron Kernel Interface (NKI), through which each lets developers use dedicated kernels optimized for its AI chip. In other words, even when using PyTorch's common APIs, hardware-specific optimization remains important when pursuing high performance. Developers considering adoption need to check the range of supported operations as well as compilation frequency, cache reuse, and actual processing times.
The latest moves by Google and AWS draw attention as a way of widening freedom of hardware choice in AI development. The software environment built around NVIDIA's CUDA has so far been a major strength, but TorchTPU and TorchNeuron are creating environments in which other AI chips can be tried while keeping widely used PyTorch code and libraries. However, ease of adoption alone does not determine which option is better in real performance or operating cost. Comparing processing speed, compilation time, the range of executable operations, and operating cost under the same models and workloads should show how strong an alternative to existing GPU environments TPUs and Trainium can become.