Migrating to JAX jax-v0.9.2
Version jax-v0.9.2 introduces 2 breaking changes. This guide details how to update your code.
Released: 3/18/2026
⚠️ Check Your Code
If you use any of these symbols, you need to read this guide:
jax._src.literals.TypedNdArrayjax.numpy.arangeBreaking Changes
●Issue #1
The type `jax._src.literals.TypedNdArray` is now a subclass of `np.ndarray` instead of a duck type. Code relying on duck typing might need updates.
●Issue #2
`jax.numpy.arange` with `step` specified no longer generates the array on host, which may lead to less precise outputs for narrow-width floats (e.g., bfloat16). To restore previous behavior for narrow-width floats, use `jnp.array(np.arange(...))`.
Migration Steps
- 1If you relied on the previous duck-typing behavior of `jax._src.literals.TypedNdArray`, update your code.
- 2If you observe precision issues with narrow-width floats (like bfloat16) when using `jnp.arange(..., step=...)`, use `jnp.array(np.arange(...))` instead.
Release Summary
JAX 0.9.2 updates the internal type structure of `TypedNdArray` and changes the execution context for `jnp.arange` with a step argument, potentially affecting float precision.
Need More Details?
View the full release notes and all changes for JAX jax-v0.9.2.
View Full Changelog