Update jax requirement from ~=0.4.38 to ~=0.9.0 in /dev_tools/requirements/deps

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.9.0

• New features:

  • Added 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. (#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.
  • jax.numpy.fix is deprecated, anticipating the deprecation of numpy.fix in NumPy v2.5.0. jax.numpy.trunc is a drop-in replacement.

• Changes:

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

Changelog

Sourced from jax's changelog.

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 deprecation warning, because the output is poorly-defined.
  • From {mod}jax.core a number of symbols are newly deprecated including: call_impl, get_aval, mapped_aval, subjaxprs, set_current_trace, take_current_trace, traverse_jaxpr_params, unmapped_aval, AbstractToken, and TraceTag.
  • All symbols in {mod}jax.interpreters.pxla are deprecated. These are primarily JAX internal APIs, and users should not rely on them.

• Changes:

  • jax's Tracer no longer inherits from jax.Array at runtime. However,

... (truncated)

Commits

• 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.
• 8f4389a Increase shard count for TPU and GPU tests back to 5 for api_test.py.
• Additional commits viewable in compare view

Most Recent Ignore Conditions Applied to This Pull Request

Dependency Name	Ignore Conditions
jax	[>= 0.5.dev0, < 0.6]

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 merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @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)