●Change8

Migrating to JAX jax-v0.9.1

Version jax-v0.9.1 introduces 1 breaking change. This guide details how to update your code.

Released: 3/2/2026

1
Breaking Changes
1
Migration Steps
2
Affected Symbols

⚠️ Check Your Code

If you use any of these symbols, you need to read this guide:

jax.shard_mapjax.reshard

Breaking Changes

●Issue #1

JAX tracers that are not of `Array` type (e.g., of `Ref` type) no longer report themselves as instances of `Array`. Code relying on this incorrect reporting will need updating.

Migration Steps

  1. 1
    If using `jax.shard_map` in Explicit mode and relying on implicit resharding when `in_specs` is provided, explicitly call `jax.reshard` on inputs before passing them to `shard_map`.

Release Summary

This release improves strictness in `jax.shard_map` explicit mode by enforcing PartitionSpec matching and updates JAX tracer type reporting. A new debug configuration for the compilation cache was also introduced.

Need More Details?

View the full release notes and all changes for JAX jax-v0.9.1.

View Full Changelog