# JAX v0.9.1 — JAX v0.9.1

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

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

---

- **changed** — JAX tracers that are not of Array type (e.g., of Ref type) will no longer report themselves to be instances of Array
- **changed** — 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, acting like an assert instead of an implicit reshard
- **added** — Added a debug config jax_compilation_cache_check_contents that verifies cache contents match when a value is put() and causes misses when get() is called on a value not put() by the current process

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