JAX v0.11.1

New features

• Added an error check for trying to deserialize JAX exports that are older than the backwards compatibility window. Without this check the deserialization of expired artifacts may succeed and then result in obscure downstream errors. Added a configuration flag --jax_export_deserialize_expired_versions to temporarily bypass the error check. See https://docs.jax.dev/en/latest/export/export.html#compatibility-guarantees.
• Added jax.numpy.top_k, which implements numpy.top_k, added in in NumPy v2.6.0 (#39729).

Breaking changes

• The exec_time_optimization_effort and memory_fitting_effort flags have been removed in favor of the EffortLevel enum.
• JAX does not support anymore deserialization of Exported modules from before January 15th, 2026 because they are beyond the backwards compatibility window. On that date we added support to serialize shardings as NamedSharding, and now that is the only sharding serialization that is supported.
• jnp.take_along_axis now always defaults wrap_negative_indices to True. It used to default to False for mode=promise_in_bounds and True otherwise. (This also means None is no longer a valid value for wrap_negative_indices.)

Deprecations

• The fields in_shardings_hlo and out_shardings_hlo of jax.export.Exported have been deprecated for a while. Now accessing them raises a warning. Use in_shardings_jax and out_shardings_jax instead.

Changes

• The cuDNN fused attention backward pass (used by jax.nn.dot_product_attention with implementation='cudnn') no longer computes a bias gradient when the only attention bias comes from a boolean mask, whose gradient no caller can request. Bias gradients for an explicit bias or a non-boolean mask are unchanged (#34685).
• jax.numpy.meshgrid, jax.numpy.ogrid, and jax.numpy.broadcast_arrays now return tuples rather than lists in order to align with NumPy>2.0 and the Array API specification. (#39783, #39789, #39802)
• When jax.grad or jax.value_and_grad rejects a function with a non-scalar output, the error message now suggests reducing the output to a scalar (e.g. with output.sum()), using jax.jacobian, or reshaping size-1 outputs (#2303).
• When indexing with non-static or traced slice indices, the error message now suggests using jax.lax.dynamic_slice, jax.lax.dynamic_update_slice, or jax.ds, and shows tracer provenance (#7222).
• PyTree metadata equality comparison failures now report the specific registered pytree type that caused the error (#13027).

Bug fixes

• jax.numpy.linalg.det and jax.numpy.linalg.slogdet now use a closed-form LU decomposition with row pivoting for 2x2 and 3x3 matrices instead of closed-form polynomial expansions to avoid numerical instability and catastrophic cancellation (#39905).
• The batching rules of the cuDNN fused attention primitives (used by jax.nn.dot_product_attention with implementation='cudnn') now support operands that do not carry the vmap axis, including a shared bias or mask. Previously jax.jacobian, jax.vmap with partial in_axes, and jax.vmap of a VJP or of jax.grad failed with a reshape TypeError (#38495).
• jax.vmap of fp8 cuDNN fused attention now works: its batching rules additionally mislabeled or dropped the amax outputs and restored output shapes incorrectly, so previously no vmap of the fp8 path succeeded at all. The amax outputs are whole-batch statistics and do not carry the vmap axis; vmap over the scale/descale operands raises a clear NotImplementedError.
• Setting jax_compiler_enable_remat_pass to False now adds rematerialization to the set of disabled XLA passes instead of overwriting it, so HLO passes disabled via XLA_FLAGS=--xla_disable_hlo_passes=... stay disabled (#37391).
• jax.numpy.split, jax.numpy.array_split, and the hsplit/vsplit/dsplit variants once again accept negative entries in indices_or_sections, resolving them against the axis size as NumPy does (#6599). Out-of-bound indices are now clipped to the axis bounds and produce empty sections, also matching NumPy, instead of raising ValueError: Sizes passed to split must be nonnegative.
• Fixed abstract evaluation in jax.lax.scan to only check .mat equivalency when the abstract value is a ShapedArray (#39700).
• Fixed propagation of singleton sharded dimensions in jax.lax.reshape when reshaping arrays with sharding constraints (#39309).
• Fixed jax.tree_util.flatten_one_level_with_keys for namedtuple instances (#39297).
• Fixed _get_prime_factors in jax.experimental.mesh_utils (#38286).