●Change8

jax-v0.11.1

Breaking Changes
📦 jaxView on GitHub →
⚠ 3 breaking✨ 2 features🐛 9 fixes⚡ 1 deprecations🔧 20 symbols

Summary

This release introduces `jax.numpy.top_k` and improves error checking for JAX exports. It also enforces stricter compatibility for deserialized exports and refines several API behaviors for better alignment with NumPy and Array API specifications.

⚠️ Breaking Changes

  • The `exec_time_optimization_effort` and `memory_fitting_effort` flags have been removed. Use the `EffortLevel` enum instead.
  • JAX no longer supports deserialization of Exported modules from before January 15th, 2026. Only `NamedSharding` serialization is supported.
  • jnp.take_along_axis now always defaults wrap_negative_indices to True. The previous default behavior varied based on the mode. `None` is no longer a valid value for wrap_negative_indices.

Migration Steps

  1. Replace `exec_time_optimization_effort` and `memory_fitting_effort` flags with the `EffortLevel` enum.
  2. Update code to use `in_shardings_jax` and `out_shardings_jax` instead of `in_shardings_hlo` and `out_shardings_hlo` for `jax.export.Exported`.
  3. Ensure `wrap_negative_indices` is explicitly set for `jnp.take_along_axis` if `False` is desired, as it now defaults to `True`.

✨ New Features

  • Added an error check for deserializing JAX exports older than the backwards compatibility window. A configuration flag `--jax_export_deserialize_expired_versions` can temporarily bypass this check.
  • Added `jax.numpy.top_k`, mirroring `numpy.top_k`.

🐛 Bug Fixes

  • `jax.numpy.linalg.det` and `jax.numpy.linalg.slogdet` now use LU decomposition for 2x2 and 3x3 matrices to improve numerical stability.
  • Batching rules for cuDNN fused attention primitives now support operands without a vmap axis, resolving `TypeError` issues with `jax.jacobian`, `jax.vmap` with partial `in_axes`, and `jax.vmap` of VJP or `jax.grad`.
  • Vmap of fp8 cuDNN fused attention now works correctly, fixing issues with amax outputs and output shapes.
  • Setting `jax_compiler_enable_remat_pass` to `False` now correctly adds `rematerialization` to disabled XLA passes without overwriting existing disabled passes.
  • `jax.numpy.split`, `jax.numpy.array_split`, and `hsplit`/`vsplit`/`dsplit` now accept negative entries in `indices_or_sections` and clip out-of-bound indices, matching NumPy's behavior.
  • Fixed abstract evaluation in `jax.lax.scan` to only check `.mat` equivalency for `ShapedArray`.
  • Fixed propagation of singleton sharded dimensions in `jax.lax.reshape` for arrays with sharding constraints.
  • Fixed `jax.tree_util.flatten_one_level_with_keys` for `namedtuple` instances.
  • Fixed `_get_prime_factors` in `jax.experimental.mesh_utils`.

Affected Symbols

jax.export.Exportedjax.numpy.top_kjnp.take_along_axisjax.numpy.meshgridjax.numpy.ogridjax.numpy.broadcast_arraysjax.gradjax.value_and_gradjax.lax.dynamic_slicejax.lax.dynamic_update_slicejax.dsjax.nn.dot_product_attentionjax.numpy.linalg.detjax.numpy.linalg.slogdetjax.jacobianjax.vmapjax.lax.scanjax.lax.reshapejax.tree_util.flatten_one_level_with_keysjax.experimental.mesh_utils

⚡ Deprecations

  • Accessing `in_shardings_hlo` and `out_shardings_hlo` of `jax.export.Exported` now raises a warning. Use `in_shardings_jax` and `out_shardings_jax` instead.