This document is relevant for: Inf2, Trn1, Trn2

JAX Neuron (beta)

JAX Neuron 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

• JAX Neuron Known issues

This document is relevant for: Inf2, Trn1, Trn2