This document is relevant for: Trn2, Trn3

JAX Setup

This guide provides step-by-step instructions for installing and configuring JAX with the NeuronX plugin on AWS Trainium and Inferentia instances. JAX NeuronX enables high-performance machine learning workloads by integrating JAX with AWS’s custom ML accelerators.

For more installation and deployment options, see Install JAX for Neuron.

Note

This setup guide is relevant for Inf2 & Trn1 / Trn1n / Trn2 instances.

Install JAX for Neuron

Go to the JAX setup hub for DLAMI, Deep Learning Container, and manual pip installation instructions.

This document is relevant for: Trn2, Trn3