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).