Update jax requirement from !=0.5.1,!=0.5.2,!=0.6.0,!=0.7.1,<0.8.0,>=0.5.0 to >=0.5.0,!=0.5.1,!=0.5.2,!=0.6.0,!=0.7.1,<0.10.0 in the jax group

[dependabot]

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

Release notes

Sourced from jax's releases.

JAX v0.9.0.1 is identical to v0.9.0 with the commits from the following four PRs patched in:

• openxla/xla#36579
• openxla/xla#36345
• openxla/xla#36755
• openxla/xla#36696

Changelog

Sourced from jax's changelog.

JAX 0.9.0.1

JAX v0.9.0.1 is identical to v0.9.0 with the commits from the following four PRs patched in:

• openxla/xla#36579
• openxla/xla#36345
• openxla/xla#36755
• openxla/xla#36696

JAX 0.9.0 (January 20, 2026)

• New features:

  • Added {func}jax.thread_guard, a context manager that detects when devices are used by multiple threads in multi-controller JAX.

• Bug fixes:

  • Fixed a workspace size calculation error for pivoted QR (magma_zgeqp3_gpu) in MAGMA 2.9.0 when using use_magma=True and pivoting=True. ({jax-issue}[#34145](https://github.com/jax-ml/jax/issues/34145)).

• Deprecations:

  • The flag jax_collectives_common_channel_id was removed.
  • The jax_pmap_no_rank_reduction config state has been removed. The no-rank-reduction behavior is now the only supported behavior: a jax.pmapped function f sees inputs of the same rank as the input to jax.pmap(f). For example, if jax.pmap(f) receives shape (8, 128) on 8 devices, then f receives shape (1, 128).
  • Setting the jax_pmap_shmap_merge config state is deprecated in JAX v0.9.0 and will be removed in JAX v0.10.0.
  • {func}jax.numpy.fix is deprecated, anticipating the deprecation of {func}numpy.fix in NumPy v2.5.0. {func}jax.numpy.trunc is a drop-in replacement.

• Changes:

  • {func}jax.export now supports explicit sharding. This required a new export serialization format version that includes the NamedSharding, including the abstract mesh, and the partition spec. As part of this change we have added a restriction in the use of exported modules: when calling them the abstract mesh must match the one used at export time, including the axis names. Previously, only the number of the devices mattered.

JAX 0.8.2 (December 18, 2025)

• Deprecations
  • jax.lax.pvary has been deprecated. Please use jax.lax.pcast(..., to='varying') as the replacement.
  • Complex arguments passed to {func}jax.numpy.arange now result in a

... (truncated)

Commits

• 470cf68 Prepare for JAX release 0.9.0.1
• 2de5b8b Use maxsize=None with trace_to_jaxpr's weakref_lru_cache to get more cache hi...
• 98d6b4b Remove failing test_itof_dot_canonicalization_fails_without_compat_mode test.
• 28799d5 Disable tpu_splash_attention_kernel_test on TPU v7x.
• 80b1ef6 Add libtpu guard to failing tpu_trace_value_test.
• 77d9ffb Skip tpu_pallas_distributed_test on 7x.
• 9e7b005 Add libtpu date guard for failing test.
• 18b34a1 Re-enable socket_transfer_test internally.
• d97eed6 Disable tests/multiprocess:socket_transfer_test since it's failing flakily ...
• fae5def Fix wheel sources tests for Windows platform.
• 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 <dependency name> major version will close this group update PR and stop Dependabot creating any more for the specific dependency's major version (unless you unignore this specific dependency's major version or upgrade to it yourself)
• @dependabot ignore <dependency name> minor version will close this group update PR and stop Dependabot creating any more for the specific dependency's minor version (unless you unignore this specific dependency's minor version or upgrade to it yourself)
• @dependabot ignore <dependency name> will close this group update PR and stop Dependabot creating any more for the specific dependency (unless you unignore this specific dependency or upgrade to it yourself)
• @dependabot unignore <dependency name> will remove all of the ignore conditions of the specified dependency
• @dependabot unignore <dependency name> <ignore condition> will remove the ignore condition of the specified dependency and ignore conditions

[github-actions]

Memory benchmark result

|               Test Name                |      %Δ      |    Master (MB)     |      PR (MB)       |    Δ (MB)    |    Time PR (s)     |  Time Master (s)   |
| -------------------------------------- | ------------ | ------------------ | ------------------ | ------------ | ------------------ | ------------------ |
  test_objective_jac_w7x                 |    5.45 %    |     3.828e+03      |     4.037e+03      |    208.80    |       39.47        |       37.37        |
  test_proximal_jac_w7x_with_eq_update   |    1.31 %    |     6.503e+03      |     6.588e+03      |    85.11     |       163.93       |       164.73       |
  test_proximal_freeb_jac                |    0.34 %    |     1.318e+04      |     1.322e+04      |    44.36     |       87.37        |       84.65        |
  test_proximal_freeb_jac_blocked        |    1.22 %    |     7.469e+03      |     7.560e+03      |    91.32     |       74.31        |       74.03        |
  test_proximal_freeb_jac_batched        |   -0.40 %    |     7.484e+03      |     7.455e+03      |    -29.70    |       74.22        |       73.93        |
  test_proximal_jac_ripple               |   -1.96 %    |     3.560e+03      |     3.490e+03      |    -69.79    |       65.60        |       65.90        |
  test_proximal_jac_ripple_bounce1d      |   -0.29 %    |     3.517e+03      |     3.507e+03      |    -10.29    |       77.16        |       77.48        |
  test_eq_solve                          |    0.68 %    |     1.977e+03      |     1.991e+03      |    13.38     |       95.21        |       95.51        |

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.52%. Comparing base (7fbae3a) to head (e099298).

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #2099   +/-   ##
=======================================
  Coverage   94.52%   94.52%           
=======================================
  Files         102      102           
  Lines       28785    28785           
=======================================
  Hits        27209    27209           
  Misses       1576     1576           

see 2 files with indirect coverage changes

🚀 New features to boost your workflow:

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