Update jax requirement from !=0.7.1,<0.10.0,>=0.6.2 to >=0.6.2,!=0.7.1,<0.11.0 in the jax group across 1 directory

[dependabot]

⚠️ Dependabot is rebasing this PR ⚠️

Rebasing might not happen immediately, so don't worry if this takes some time.

Note: if you make any changes to this PR yourself, they will take precedence over the rebase.

---

Updates the requirements on jax to permit the latest version.
Updates jax to 0.10.1

Release notes

Sourced from jax's releases.

JAX v0.10.1

• New features

  • Added ResizeMethod.AREA to jax.image.resize, which matches TensorFlow's AREA resizing (#20098).
  • Added jax.scipy.linalg.hadamard for constructing Hadamard matrices (#10144).
  • Added jax.scipy.linalg.circulant for constructing circulant matrices (#10144).
  • Added jax.scipy.linalg.dft for constructing discrete Fourier transform matrices (#10144).
  • Added jax.scipy.linalg.leslie for constructing Leslie matrices (#10144).
  • Added jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients (#10144).
  • Added jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices (#10144).
  • Added jax.scipy.linalg.helmert for constructing Helmert matrices (#10144).
  • Moved RNG APIs from "implementations" to dtypes (#27854):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

• Breaking changes

  • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

• Deprecations

  • Passing the copy, order, and ndmin arguments to jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
  • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future version of JAX, this will become an error, if you depend on using them as leaves, pass is_leaf to jax.tree.* methods.

Changelog

Sourced from jax's changelog.

JAX 0.10.1 (May 20, 2026)

• New features

  • Added ResizeMethod.AREA to {func}jax.image.resize, which matches TensorFlow's AREA resizing ({jax-issue}[#20098](https://github.com/jax-ml/jax/issues/20098)).
  • Added {func}jax.scipy.linalg.hadamard for constructing Hadamard matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.circulant for constructing circulant matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.dft for constructing discrete Fourier transform matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.leslie for constructing Leslie matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.linalg.helmert for constructing Helmert matrices ({jax-issue}[#10144](https://github.com/jax-ml/jax/issues/10144)).
  • Added {func}jax.scipy.special.boxcox and {func}jax.scipy.special.boxcox1p for the Box-Cox power transformation.
  • Moved RNG APIs from "implementations" to dtypes ({jax-issue}[#27854](https://github.com/jax-ml/jax/issues/27854)):
    • Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name.
    • jax.random.key and wrap_key_data now accept a dtype argument.

• Breaking changes

  • with mesh: context manager has been deprecated. Please use with jax.set_mesh(mesh): instead.

• Deprecations

  • Passing the copy, order, and ndmin arguments to {func}jax.numpy.array positionally is deprecated. Use keyword arguments instead. This matches the signature of numpy.array.
  • Python dict_values, generators, zip return type and iterators generally are deprecated by default when used as leaves in pytrees. In a future version of JAX, this will become an error, if you depend on using them as leaves, pass is_leaf to jax.tree.* methods.

JAX 0.10.0 (April 16, 2026)

• New features:
  • Added ResizeMethod.CUBIC_PYTORCH to {func}jax.image.resize to match PyTorch's bicubic resize ({jax-issue}[#15768](https://github.com/jax-ml/jax/issues/15768)).
  • We now support differentiation of {func}jax.lax.linalg.qr for wide matrices and when full_matrices is True.
  • LAPACK operations are now parallelized along the batch dimension on CPU.
  • Added perturb_singular argument to {func}jax.lax.linalg.tridiagonal_solve to handle singular matrices by perturbing near-zero pivots in the LU decomposition. This is useful for

... (truncated)

Commits

• 619764c Apply patch to revert XLA PR #42396 on release/0.10.1 branch
• 433d98d Prepare for JAX release 0.10.1
• 8c38c32 [Pallas] Enable disabled TPU Interpret Mode test of example kernel
• 6820cf6 [pallas:sc] pl.kernel is now the only API for writing Pallas SC kernels
• 456b22d [state] Fixed a latent bug in _call_primitive_discharge_rule
• d154b80 Update XLA dependency to use revision http://github.com/openxla/xla/commit/9b...
• 11af068 Update XLA dependency to use revision http://github.com/openxla/xla/commit/31...
• 1ecdea1 Improve compilation throughput in multi-threaded auto-tuning environments. Th...
• 46f75a2 Make JaxprEqnContext a bit easier to work with making it easier to add new co...
• 7c01f9b [Pallas MPMD] Refactor tracing context
• Additional commits viewable in compare view

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    5.96 %    |     3.962e+03      |     4.198e+03      |    236.18    |       32.15        |       29.74        |
  test_proximal_jac_w7x_with_eq_update   |   -0.85 %    |     6.846e+03      |     6.787e+03      |    -58.22    |       153.57       |       152.63       |
  test_proximal_freeb_jac                |   -0.03 %    |     1.345e+04      |     1.344e+04      |    -4.39     |       83.85        |       79.93        |
  test_proximal_freeb_jac_blocked        |   -0.14 %    |     7.778e+03      |     7.768e+03      |    -10.52    |       72.60        |       72.29        |
  test_proximal_freeb_jac_batched        |    0.24 %    |     7.714e+03      |     7.733e+03      |    18.40     |       72.00        |       72.30        |
  test_proximal_jac_ripple               |   -0.43 %    |     3.632e+03      |     3.617e+03      |    -15.54    |       56.26        |       56.16        |
  test_proximal_jac_ripple_bounce1d      |   -1.22 %    |     3.887e+03      |     3.840e+03      |    -47.37    |       70.62        |       70.03        |
  test_eq_solve                          |    0.63 %    |     2.093e+03      |     2.106e+03      |    13.22     |       89.45        |       89.43        |

For the memory plots, go to the summary of Memory Benchmarks workflow and download the artifact.