# JAX v0.11.1 — JAX v0.11.1

- Product: JAX (https://whatsnew.fyi/product/jax)
- Vendor: Google
- Date: 2026-08-17
- Version: v0.11.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.11.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.11.1

What's New is an index, not a publisher: every entry below links to the vendor's own release notes, which are the authoritative source. Entries are labelled where they are hand-curated sample data, pre-releases, or drawn from a secondary source such as a developer blog.

Reuse: the summaries, labels and curation here are © What's New. Quote freely with attribution and a link back; wholesale republication of the corpus is not permitted — terms: https://whatsnew.fyi/terms. The vendors' own release notes remain their publishers'.

---

- **added** — Added an error check for trying to deserialize JAX exports that are older than the backwards compatibility window, with a configuration flag `--jax_export_deserialize_expired_versions` to temporarily bypass the error check
- **added** — Added `jax.numpy.top_k`, which implements `numpy.top_k` added in NumPy v2.6.0
- **removed** — Removed the `exec_time_optimization_effort` and `memory_fitting_effort` flags in favor of the `EffortLevel` enum
- **removed** — JAX no longer supports deserialization of Exported modules from before January 15th, 2026 because they are beyond the backwards compatibility window
- **changed** — jnp.take_along_axis now always defaults wrap_negative_indices to True instead of defaulting to False for mode=promise_in_bounds and True otherwise
- **deprecated** — The fields `in_shardings_hlo` and `out_shardings_hlo` of `jax.export.Exported` now raise a warning when accessed; use `in_shardings_jax` and `out_shardings_jax` instead
- **changed** — 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`
- **changed** — `jax.numpy.meshgrid`, `jax.numpy.ogrid`, and `jax.numpy.broadcast_arrays` now return tuples rather than lists to align with NumPy>2.0 and the Array API specification
- **changed** — 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, using `jax.jacobian`, or reshaping size-1 outputs
- **changed** — 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
- **changed** — PyTree metadata equality comparison failures now report the specific registered pytree type that caused the error
- **fixed** — `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
- **fixed** — The batching rules of the cuDNN fused attention primitives now support operands that do not carry the vmap axis, including a shared bias or `mask`
- **fixed** — `jax.vmap` of fp8 cuDNN fused attention now works with corrected amax output handling and output shapes
- **fixed** — Setting `jax_compiler_enable_remat_pass` to `False` now adds `rematerialization` to the set of disabled XLA passes instead of overwriting it
- **fixed** — `jax.numpy.split`, `jax.numpy.array_split`, and the `hsplit`/`vsplit`/`dsplit` variants now accept negative entries in `indices_or_sections` and resolve them against the axis size as NumPy does
- **fixed** — Fixed abstract evaluation in `jax.lax.scan` to only check `.mat` equivalency when the abstract value is a `ShapedArray`
- **fixed** — Fixed propagation of singleton sharded dimensions in `jax.lax.reshape` when reshaping arrays with sharding constraints
- **fixed** — Fixed `jax.tree_util.flatten_one_level_with_keys` for `namedtuple` instances
- **fixed** — Fixed `_get_prime_factors` in `jax.experimental.mesh_utils`

* 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_FL

_[Truncated at 4000 characters — full notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.11.1]_