# JAX v0.9.0 — JAX v0.9.0

- Product: JAX (https://whatsnew.fyi/product/jax)
- Vendor: Google
- Date: 2026-01-20
- Version: v0.9.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.9.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.9.0

What's New is an index, not a publisher: every entry below links to the vendor's own release notes, which are the authoritative source. Entries are labelled where they are hand-curated sample data, pre-releases, or drawn from a secondary source such as a developer blog.

Reuse: the summaries, labels and curation here are © What's New. Quote freely with attribution and a link back; wholesale republication of the corpus is not permitted — terms: https://whatsnew.fyi/terms. The vendors' own release notes remain their publishers'.

---

- **added** — Added jax.thread_guard, a context manager that detects when devices are used by multiple threads in multi-controller JAX
- **fixed** — 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
- **removed** — Removed the flag jax_collectives_common_channel_id
- **removed** — Removed the jax_pmap_no_rank_reduction config state; the no-rank-reduction behavior is now the only supported behavior
- **deprecated** — Setting the jax_pmap_shmap_merge config state is deprecated and will be removed in JAX v0.10.0
- **deprecated** — jax.numpy.fix is deprecated; jax.numpy.trunc is a drop-in replacement
- **changed** — jax.export now supports explicit sharding with a new export serialization format version that includes NamedSharding, the abstract mesh, and the partition spec
- **changed** — When calling exported modules, the abstract mesh must now match the one used at export time, including the axis names

* 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.pmap`ped 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.