This document is relevant for: Inf1, Inf2, Trn1, Trn1n

JAX Neuron (beta)

JAX NeuronX is a software package containing Neuron-specific JAX features, such as the Neuron NKI JAX interface. It also serves as a meta-package for providing a tested combination of the jax-neuronx, jax, jaxlib, libneuronxla , and neuronx-cc packages.

• JAX Neuron plugin Setup

This document is relevant for: Inf1, Inf2, Trn1, Trn1n