This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3
Install JAX for Neuron#
Install JAX with AWS Neuron support for training and inference on Inferentia and Trainium instances.
Supported Instances: Inf2, Trn1, Trn2, Trn3
JAX Version: 0.10+ with Neuron PJRT plugin
Beta Release
JAX NeuronX is currently in beta. Some JAX functionality may not be fully supported. We welcome your feedback and contributions.
Choose Installation Method#
Prerequisites#
Before installing, ensure you have:
Requirement |
Details |
|---|---|
Instance Type |
Inf2, Trn1, Trn2, or Trn3 instance |
Operating System |
Ubuntu 26.04, Ubuntu 24.04, or Amazon Linux 2023. For the tested Linux kernel and GLIBC versions, see Neuron Runtime and Driver compatibility reference. |
Python Version |
Python 3.11, 3.12, 3.13, or 3.14 |
AWS Account |
With EC2 launch permissions |
SSH Access |
Key pair for instance connection |
What You’ll Get#
After installation, you’ll have:
JAX 0.10+ with Neuron PJRT plugin
jax-neuronx package for Neuron-specific features
libneuronxla PJRT plugin for native JAX device integration
neuronx-cc compiler for model optimization
Neuron Runtime for model execution
Version Information#
Component |
Version |
|---|---|
JAX |
0.10.0+ |
jax-neuronx |
0.10.0+ |
libneuronxla |
latest |
neuronx-cc |
2.28.0+ |
Python |
3.11, 3.12, 3.13, 3.14 |
JAX and jaxlib version compatibility#
The jax-neuronx[stable] meta-package is the recommended way to install JAX for Neuron. It pins a tested combination of jax, jaxlib, libneuronxla, and neuronx-cc, so you get a known-good environment without choosing versions yourself. The DLAMI, DLC, and manual installation methods above all use it.
If you need to bring your own jax and jaxlib versions, install the packages separately instead. Because libneuronxla supports a broad range of jaxlib versions through the PJRT C-API mechanism, this gives you flexibility when choosing jax and jaxlib versions and lets you add the Neuron plugin to an existing JAX environment:
python3 -m pip install jax==0.10.0 jaxlib==0.10.0
python3 -m pip install jax-neuronx libneuronxla neuronx-cc==2.* --extra-index-url=https://pip.repos.neuron.amazonaws.com
Replace 0.10.0 with the jax and jaxlib version tested for your Neuron release. See the JAX NeuronX release notes for the version shipped in each release.
To check which client versions your installed libneuronxla supports, use the libneuronxla.supported_clients API:
import libneuronxla
libneuronxla.supported_clients()
It returns a list of strings in the form "<package> <version> (PJRT C-API <c-api version>)" — one entry per known-supported jaxlib build. Run it in your own environment to see the exact versions for your installed libneuronxla.
The list covers known versions only and may be incomplete. Other versions can work, including future jaxlib releases, as long as the PJRT C-API stays compatible with your installed libneuronxla. For this reason, Neuron does not pin a strict libneuronxla –> jaxlib dependency, which gives you more freedom when coordinating jax and libneuronxla installations. Compatibility between jaxlib and libneuronxla is determined by the PJRT C-API version. For more information, see the PJRT integration guide.
Next Steps#
After installation:
Verify Installation: Run verification commands in the installation guide
Explore JAX on Neuron: JAX Support on Neuron
API Reference: API Reference Guide for JAX Neuronx
This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3