●Change8

Migrating to JAX jax-v0.8.0

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

Released: 10/15/2025

19
Breaking Changes
12
Migration Steps
36
Affected Symbols

⚠️ Check Your Code

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

jax.pmapjax.experimental.shard_map.shard_mapjax.numpy.covjax.interpreters.mlir.custom_calljax.utiljax.extend.ffijax.experimental.host_callbackjax.custom_derivatives.custom_jvp_call_jaxpr_pjax.experimental.multihost_utils.process_allgatherjax.experimental.compilation_cache.is_initializedjax.experimental.compilation_cache.initialize_cachejax.interpreters.xla.canonicalize_dtypejaxlib.hlo_helpersjax_cpu_enable_gloo_collectivesjax.numpy.percentilejax.numpy.quantilefor_loopjax.numpy.trimzerosjax.numpy.sumjax.dlpackjax.errorsjax.lib.xla_bridgejax.lib.xla_clientjax.lib.xla_extensionjax.interpreters.mlir.dense_bool_arrayjax.numpy.linalg.eigjax.gradjax.vjpjax.dlpack.from_dlpackjax.lax.linalg.eigjax.numpy.trim_zerosjax.enable_x64jax.experimental.enable_x64jax.experimental.disable_x64jax.experimental.pjit.pjitjax.jit

Breaking Changes

●Issue #1

Default implementation of `jax.pmap` switched to use `jax.jit` and `jax.shard_map`; new code should call `jax.shard_map` directly (see migration guide).

●Issue #2

`auto=` parameter removed from `jax.experimental.shard_map.shard_map`; nesting no longer supported – use `jax.shard_map` for nested calls.

●Issue #3

Objects implementing `__jax_array__` can no longer be passed directly to JIT‑compiled functions; wrap them with `jax.numpy.asarray` first.

●Issue #4

`jax.numpy.cov` now returns NaN for empty arrays and matches NumPy 2.2 behavior for single‑row design matrices.

●Issue #5

`Array` values are no longer accepted where a `dtype` is expected; extract the dtype via `.dtype` before passing.

●Issue #6

Removed deprecated function `jax.interpreters.mlir.custom_call`.

●Issue #7

Modules `jax.util`, `jax.extend.ffi`, and `jax.experimental.host_callback` have been removed.

●Issue #8

Removed deprecated symbol `jax.custom_derivatives.custom_jvp_call_jaxpr_p`.

●Issue #9

`jax.experimental.multihost_utils.process_allgather` now errors when given a non‑fully‑addressable `jax.Array` with `tiled=False`; pass `tiled=True`.

●Issue #10

Deprecated symbols `is_initialized` and `initialize_cache` removed from `jax.experimental.compilation_cache`.

●Issue #11

Removed deprecated function `jax.interpreters.xla.canonicalize_dtype`.

●Issue #12

`jaxlib.hlo_helpers` removed; use `jax.ffi` instead.

●Issue #13

Option `jax_cpu_enable_gloo_collectives` removed; use `jax_cpu_collectives_implementation`.

●Issue #14

`interpolation` argument removed from `jax.numpy.percentile` and `jax.numpy.quantile`; use `method` argument.

●Issue #15

Internal `for_loop` primitive removed; use `jax.lax.fori_loop` directly.

●Issue #16

`jax.numpy.trimzeros` now raises an error for non‑1D input.

●Issue #17

`where` argument to reduction functions like `jax.numpy.sum` must now be boolean.

●Issue #18

Removed deprecated functions in `jax.dlpack`, `jax.errors`, `jax.lib.xla_bridge`, `jax.lib.xla_client`, and `jax.lib.xla_extension`.

●Issue #19

Removed `jax.interpreters.mlir.dense_bool_array`; construct attributes via MLIR APIs.

Migration Steps

  1. 1
    Update code to use `jax.shard_map` instead of `jax.pmap` or the deprecated `jax.experimental.shard_map.shard_map`.
  2. 2
    If you relied on nesting `shard_map`, switch to `jax.shard_map` for nested calls.
  3. 3
    Wrap objects with `__jax_array__` using `jax.numpy.asarray` before passing to JIT‑compiled functions.
  4. 4
    When calling `jax.experimental.multihost_utils.process_allgather` with a `jax.Array`, add `tiled=True`.
  5. 5
    Replace imports from removed modules (`jax.util`, `jax.extend.ffi`, `jax.experimental.host_callback`) with appropriate alternatives (e.g., `jax.ffi`).
  6. 6
    Use `jax.ffi` instead of the removed `jaxlib.hlo_helpers`.
  7. 7
    Switch from the removed `jax_cpu_enable_gloo_collectives` option to `jax_cpu_collectives_implementation`.
  8. 8
    Change calls to `jax.numpy.percentile` and `jax.numpy.quantile` to use the `method` argument rather than `interpolation`.
  9. 9
    Replace any use of the internal `for_loop` primitive with `jax.lax.fori_loop`.
  10. 10
    Ensure the `where` argument passed to reduction ops like `jax.numpy.sum` is a boolean array.
  11. 11
    Update any code that used the removed deprecated functions in `jax.dlpack`, `jax.errors`, `jax.lib.xla_*` modules to the current public APIs.
  12. 12
    If you used `jax.interpreters.mlir.dense_bool_array`, construct MLIR attributes via the MLIR APIs instead.

Release Summary

JAX introduces several breaking changes, including a new default implementation for `jax.pmap` and removal of many deprecated APIs, while adding new features such as namedtuple returns for eig and enhanced dlpack support.

Need More Details?

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

View Full Changelog