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

Announcing end of support for jax_neuronx.nki_call API in jax-neuronx from starting next release

Starting with Neuron Release 2.23, Neuron will end support for 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 are recommended to switch invocations to directly call functions annotated with nki.jit API.

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