.. meta::
   :description: Install and set up JAX NeuronX plugin for AWS Trainium and Inferentia instances. Complete setup guide for JAX on Ubuntu 22 with Neuron SDK integration.

.. _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.

.. contents:: Table of contents
   :local:
   :depth: 2


``JAX`` setup on Ubuntu 22
---------------------------

.. toctree::
   :maxdepth: 1
   :hidden:

   JAX NeuronX on Ubuntu 22 </frameworks/jax/setup/jax-setup>

.. card:: Ubuntu 22 (Ubuntu22 AMI)
        :link: /frameworks/jax/setup/jax-setup
        :link-type: doc
        :class-body: sphinx-design-class-title-small