Bump jax from 0.9.0.1 to 0.9.2

[dependabot]

Bumps jax from 0.9.0.1 to 0.9.2.

Release notes

Sourced from jax's releases.

JAX v0.9.2

JAX 0.9.2 (March 2, 2026)

• Changes:
  • The semi-private type jax._src.literals.TypedNdArray is now a subclass of np.ndarray, rather than a duck type of it.
  • jax.numpy.arange with step specified no longer generates the array on host. The benefit is more efficient code, though this can lead to less precise outputs for narrow-width floats (e.g. bfloat16). To recover the previous behavior in this case, use jnp.array(np.arange(...)).

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.2 (March 18, 2026)

• Changes:
  • The semi-private type jax._src.literals.TypedNdArray is now a subclass of np.ndarray, rather than a duck type of it.
  • {func}jax.numpy.arange with step specified no longer generates the array on host. The benefit is more efficient code, though this can lead to less precise outputs for narrow-width floats (e.g. bfloat16). To recover the previous behavior in this case, use jnp.array(np.arange(...)).

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

... (truncated)

Commits

• a659757 Prepare for JAX release 0.9.2
• e0bb88a Cherry-pick upstream XLA and Shardy fixes via Bazel patches.
• 43195de Guard tests that fail in libtpu==0.0.37, but succeed at nightly libtpu==0.0.3...
• 6ca5383 Make sure to use to_ct_spec() when unsharding Zeros in shard_map_transpose....
• ed8e9d5 [Mosaic GPU] Add Pallas lowering for AsyncStoreSmemToTmem op.
• a86c180 [Mosaic GPU] Make tests involving global state order-independent.
• 7f265b4 [NFC] Reduce nesting with early return in _core_map_lowering_rule.
• e825397 Update XLA dependency to use revision http://github.com/openxla/xla/commit/18...
• 8304fec Update XLA dependency to use revision http://github.com/openxla/xla/commit/52...
• 13dcc92 Merge pull request #34332 from jakevdp:searchsorted-hijax
• 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)

[github-actions]

Thank you for your contribution @dependabot[bot] 🚀! Your github-pages is ready for download 👉 here 👈!
_{(The artifact expires on 2026-04-13T06:05:14Z. You can re-generate it by re-running the workflow here.)}