JAX v0.9.2

JAX 0.9.2 (March 2, 2026)

• Changes:
  • The semi-private type jax._src.literals.TypedNdArray is now a subclass of np.ndarray, rather than a duck type of it.
  • jax.numpy.arange with step specified no longer generates the array on host. The benefit is more efficient code, though this can lead to less precise outputs for narrow-width floats (e.g. bfloat16). To recover the previous behavior in this case, use jnp.array(np.arange(...)).