JAX NeuronX (jax-neuronx) release notes — AWS Neuron Documentation (original) (raw)

Contents

• Release [0.5.3.1.0.0]
  • Summary
• Release [0.1.3]
  • Summary
• Release [0.1.2]
  • Summary
  • What’s new in this release
• Release [0.1.1]
  • Summary
  • What’s new in this release

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

JAX NeuronX (jax-neuronx) release notes#

Table of Contents

• Release [0.5.3.1.0.0]
• Release [0.1.3]
• Release [0.1.2]
• Release [0.1.1]

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#

• This release supports JAX versions up to 0.5.0.
• Known issues are listed within JAX NeuronX Known Issues.

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