●Change8

Migrating to JAX jax-v0.6.1

Version jax-v0.6.1 introduces 3 breaking changes. This guide details how to update your code.

Released: 5/21/2025

3
Breaking Changes
3
Migration Steps
4
Affected Symbols

⚠️ Check Your Code

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

jax.lax.axis_sizejax.sharding.PartitionSpecjax.ShapeDtypeStructjax.custom_derivatives.custom_jvp_call_jaxpr_p

Breaking Changes

●Issue #1

`jax.sharding.PartitionSpec` no longer inherits from a tuple, breaking code that relied on tuple behavior. Fix by treating it as a regular object and accessing its fields directly.

●Issue #2

`jax.ShapeDtypeStruct` is now immutable; in‑place updates will raise errors. Use the `.update` method to create a modified copy.

●Issue #3

Re‑enabled strict CUDA dependency version checks may cause installation failures if incompatible CUDA packages are present. Ensure your CUDA packages meet the required versions.

Migration Steps

  1. 1
    Update any code that treats `jax.sharding.PartitionSpec` as a tuple to use its explicit attributes.
  2. 2
    Replace in‑place modifications of `jax.ShapeDtypeStruct` with calls to its `.update` method.
  3. 3
    Verify that installed CUDA packages satisfy the version requirements now enforced by JAX.

Release Summary

This release adds the new `jax.lax.axis_size` feature, makes `PartitionSpec` and `ShapeDtypeStruct` behavior stricter, re‑enables CUDA version checks, and deprecates `custom_jvp_call_jaxpr_p`.

Need More Details?

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

View Full Changelog