Bump jax from 0.9.0.1 to 0.9.1

[dependabot]

Bumps jax from 0.9.0.1 to 0.9.1.

Release notes

Sourced from jax's releases.

JAX v0.9.1

• Changes:

  • JAX tracers that are not of Array type (e.g., of Ref type) will no longer report themselves to be instances of Array.
  • Using jax.shard_map in Explicit mode will raise an error if the PartitionSpec of input does not match the PartitionSpec specified in in_specs. In other words, it will act like an assert instead of an implicit reshard. in_specs is an optional argument so you can omit specifying it and shard_map will infer the PartitionSpec from the argument. If you want to reshard your inputs, you can use jax.reshard on the arguments and then pass those args to shard_map.

• New features:

  • Added a debug config jax_compilation_cache_check_contents. If set, we miss when get() is called on a value that has not been put() by the current process, even if the value is actually in the disk cache. When a value is put(), we verify that its contents match.

Changelog

Sourced from jax's changelog.

JAX 0.9.1 (March 2, 2026)

• Changes:

  • JAX tracers that are not of Array type (e.g., of Ref type) will no longer report themselves to be instances of Array.
  • Using jax.shard_map in Explicit mode will raise an error if the PartitionSpec of input does not match the PartitionSpec specified in in_specs. In other words, it will act like an assert instead of an implicit reshard. in_specs is an optional argument so you can omit specifying it and shard_map will infer the PartitionSpec from the argument. If you want to reshard your inputs, you can use jax.reshard on the arguments and then pass those args to shard_map.

• New features:

  • Added a debug config jax_compilation_cache_check_contents. If set, we miss when get() is called on a value that has not been put() by the current process, even if the value is actually in the disk cache. When a value is put(), we verify that its contents match.

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

... (truncated)

Commits

• 58cb6e5 Prepare for JAX release 0.9.1
• bb3362e Skip sparse random and dot general sampled ad tests.
• 7aa10a0 Skip layout_test.test_layout_donation_mismatching_in_and_out_fails for TPU v7...
• 162b2ed Fix release workflow by using the folder with pre-downloaded jax wheels as ar...
• 70a8149 Rename pypi_latest libtpu-version-type to pinned and use pinned libTPU ve...
• 0eb2314 Skip tests that are failing at libtpu=0.0.35, but passing at libtpu-0.0.37.de...
• 969da1b [linux-arm64] Reduce parallelism to 32 for nogil pytest cpu to avoid OOMs.
• a573e86 Skip pallas lowering determinism testOrderAgnostic test if not on TPU.
• 6428a91 [jaxlib] Ported normalize_stubs.sh to Python
• 17cfbb8 Add jaxlib version guard for updated make_c_api_client.
• 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)