Update jax requirement from <0.11,>=0.10.2 to >=0.10.2,<0.12

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

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

... (truncated)

Changelog

Sourced from jax's changelog.

JAX 0.11.1 (August 17, 2026)

• 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 {func}jax.numpy.top_k, which implements {func}numpy.top_k, added in in NumPy v2.6.0 ({jax-issue}[#39729](https://github.com/jax-ml/jax/issues/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 {func}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 ({jax-issue}[#34685](https://github.com/jax-ml/jax/issues/34685)).
  • {func}jax.numpy.meshgrid, {obj}jax.numpy.ogrid, and {func}jax.numpy.broadcast_arrays now return tuples rather than lists in order to align with NumPy>2.0 and the Array API specification. ({jax-issue}[#39783](https://github.com/jax-ml/jax/issues/39783), {jax-issue}[#39789](https://github.com/jax-ml/jax/issues/39789), {jax-issue}[#39802](https://github.com/jax-ml/jax/issues/39802))
  • When {func}jax.grad or {func}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 {func}jax.jacobian, or reshaping size-1 outputs ({jax-issue}[#2303](https://github.com/jax-ml/jax/issues/2303)).
  • When indexing with non-static or traced slice indices, the error message now suggests using {func}jax.lax.dynamic_slice, {func}jax.lax.dynamic_update_slice, or jax.ds, and shows tracer provenance ({jax-issue}[#7222](https://github.com/jax-ml/jax/issues/7222)).
  • PyTree metadata equality comparison failures now report the specific registered pytree type that caused the error ({jax-issue}[#13027](https://github.com/jax-ml/jax/issues/13027)).

... (truncated)

Commits

• 2d66622 Prepare for JAX release 0.11.1
• 3ba3d76 [XLA:GPU] Update lax test precision to account for numerics of new cuDNN fron...
• 366ff35 Update XLA dependency to use revision
• ee8d0a1 Rename BUILD to BUILD.bazel in JAX toolchains.
• e309df5 Merge pull request #40032 from hawkinsp:bazel
• 36cfce1 Merge pull request #38946 from aybchan:aybchan/update-thor-gpuinfo-test
• 3b1573a Update bazelversion to 7.7.1.
• 91eac72 Merge pull request #40031 from jax-ml:fixeditable
• 2dc5e1a Removed core_map_p and everything which depended on it
• 7be361e Merge pull request #39759 from jakevdp:util-implicit-typevar
• Additional commits viewable in compare view

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.26%. Comparing base (7aa31b9) to head (d6036ce).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #106   +/-   ##
=======================================
  Coverage   85.26%   85.26%           
=======================================
  Files          14       14           
  Lines         638      638           
=======================================
  Hits          544      544           
  Misses         94       94           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.