●Change8

Migrating to JAX jax-v0.7.1

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

Released: 8/20/2025

3
Breaking Changes
5
Migration Steps
9
Affected Symbols

⚠️ Check Your Code

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

jax.set_meshjax.sharding.use_meshjax.lax.dotjax.lax.zeros_like_arrayjax.numpy.zeros_likejax.experimental.host_callbackjax.interpreters.adjax.interpreters.batchingjax.interpreters.partial_eval

Breaking Changes

●Issue #1

Removed `jax.sharding.use_mesh`; code using it will break. Replace with `jax.set_mesh`.

●Issue #2

Importing `jax.experimental.host_callback` will raise ImportError starting in JAX v0.8.0. Update imports to avoid this module.

●Issue #3

Positional use of `precision` and `preferred_element_type` in `jax.lax.dot` is deprecated and will be removed; pass them as keyword arguments.

Migration Steps

  1. 1
    Replace any usage of `jax.sharding.use_mesh` with `jax.set_mesh`.
  2. 2
    Update imports to stop using `jax.experimental.host_callback`; remove or replace its usage before upgrading to v0.8.0.
  3. 3
    Change calls to `jax.lax.dot` that pass `precision` or `preferred_element_type` positionally to use keyword arguments.
  4. 4
    Replace `jax.lax.zeros_like_array` with `jax.numpy.zeros_like`.
  5. 5
    If your code relied on internal APIs from `jax.interpreters.ad`, `jax.interpreters.batching`, or `jax.interpreters.partial_eval`, refactor to avoid them as they are deprecated.

Release Summary

JAX introduces new Python 3.13t/3.14t wheels, a new `jax.set_mesh` API, CUDA 12.9 builds, and several deprecations including removal of `jax.sharding.use_mesh`.

Need More Details?

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

View Full Changelog