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.7+ 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#

🚀 AWS Deep Learning AMI

Recommended for most users

Pre-configured environment with all dependencies

✅ All dependencies included

✅ Tested configurations

✅ Multiple Python versions

⏱️ Setup time: ~5 minutes

🐳 Deep Learning Container

For containerized deployments

Pre-configured Docker images from AWS ECR

✅ Docker-based isolation

✅ Training and inference images

✅ Training images available

⏱️ Setup time: ~10 minutes

🔧 Manual Installation

For custom environments

Install on existing systems or custom setups

✅ Existing system integration

✅ Custom Python versions

✅ Full control over dependencies

⏱️ Setup time: ~15 minutes

Before installing, ensure you have:

Requirement	Details
Instance Type	Inf2, Trn1, Trn2, or Trn3 instance
Operating System	Ubuntu 24.04, Ubuntu 22.04, or Amazon Linux 2023
Python Version	Python 3.10, 3.11, or 3.12
AWS Account	With EC2 launch permissions
SSH Access	Key pair for instance connection

After installation, you’ll have:

After installation:

This document is relevant for: Inf1, Inf2, Trn1, Trn2, Trn3