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

Neuron no longer supports jax_neuronx.nki_call API in jax-neuronx starting this release

Neuron Release 2.23 no longer supports jax_neuronx.nki_call API in jax-neuronx package.

For a full list of features that require jax-neuronx, please see jax_neuron-known-issues.

Customers using jax_neuronx.nki_call API will need to switch invocations to directly call functions annotated with nki.jit API.

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