Yes. AWS's Neuron documentation lists native PyTorch support on Trainium2 and Trainium3 through TorchNeuron, which it labels Beta.
What the Neuron docs say
AWS Neuron documentation says TorchNeuron supports eager mode, torch.compile with backend="neuron" and standard distributed APIs, and that you can change the device "from cuda to neuron." Earlier components such as torch-neuronx are described as being for maintaining existing workloads only.
AWS's Trainium page goes further, saying PyTorch, Hugging Face and Ray "work without modification." Treat that as AWS's claim and test your own model, since the Neuron docs mark the native path as Beta.
JAX is also available on Neuron as a Beta integration. For the SDK itself, see what the Neuron SDK is, and for the chip, what AWS Trainium is. The wider landscape is in NVIDIA's main competitors in AI.
