JAX Neuron (beta) — AWS Neuron Documentation (original) (raw)

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 NeuronX Known Issues
• API Reference Guide for JAX Neuronx
• JAX NeuronX (jax-neuronx) release notes

This document is relevant for: Inf2, Trn1, Trn2