Jax: jax-v0.9.1 Release

Release date:
February 28, 2026
Previous version:
jax-v0.9.0.1 (released February 3, 2026)
Magnitude:
17,997 Diff Delta
Contributors:
35 total committers
Data confidence:
Commits:

254 Commits in this Release

Ordered by the degree to which they evolved the repo in this version.

Authored January 19, 2026
Authored February 26, 2026
Authored February 12, 2026
Authored February 17, 2026
Authored February 18, 2026
Authored January 28, 2026

Top Contributors in jax-v0.9.1

jakevdp
bchetioui
danielsuo
allanrenucci
a-googler
emilyfertig
olupton
PatriosTheGreat
yueshengys
WindQAQ

Directory Browser for jax-v0.9.1

We haven't yet finished calculating and confirming the files and directories changed in this release. Please check back soon.

Release Notes Published

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