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

JAX NeuronX (jax-neuronx) release notes#

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.

Release [0.5.3.1.0.0]#

Date: 05/xx/2025

Summary#

  • This release supports JAX versions up to 0.5.3.

  • jax_neuronx.nki_call is no longer supported. Use neuronxcc.nki.jit instead.

  • Known issues are listed within JAX NeuronX Known Issues.

Release [0.1.3]#

Date: 04/03/2025

Summary#

Release [0.1.2]#

Date: 12/20/2024

Summary#

This release supports JAX versions up to 0.4.35.

What’s new in this release#

  • Support for JAX versions up to 0.4.35.

  • Support for JAX caching API for versions 0.4.30+.

Release [0.1.1]#

Date: 09/16/2024

Summary#

This is the initial beta release of JAX NeuronX that contains Neuron-specific JAX features, such as the Neuron NKI JAX interface

What’s new in this release#

Announcing the first JAX NeuronX release

  • JAX interface for Neuron NKI

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