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 NeuronX plugin Setup

• JAX NeuronX Known Issues

• API Reference Guide for JAX Neuronx

• JAX NeuronX (jax-neuronx) release notes

This document is relevant for: Inf2, Trn1, Trn2