Bump jaxlib from 0.4.38 to 0.9.0.1

[dependabot]

Bumps jaxlib from 0.4.38 to 0.9.0.1.

Changelog

Sourced from jaxlib'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 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)