.. meta::
   :description: Install and set up the JAX NeuronX plugin for AWS Trainium and Inferentia instances with the Neuron SDK.

.. _setup-jax-neuronx:

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 :ref:`jax-neuron-setup`.

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

.. card:: Install JAX for Neuron
        :link: jax-setup
        :link-type: ref
        :class-body: sphinx-design-class-title-small

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