●Change8

Migrating to JAX jax-v0.6.0

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

Released: 4/17/2025

5
Breaking Changes
7
Migration Steps
43
Affected Symbols

⚠️ Check Your Code

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

jax.numpy.arrayconfig.jax_data_dependent_tracing_fallbackconfig.jax_eager_pmapjax.jitjax.pmapjax.jit.lowerjax.jit.tracejax.pmap.lowerjax.pmap.tracejax.lib.xla_extensionjax.interpreters.mlir.hlojax.interpreters.mlir.func_dialectjax.interpreters.mlir.custom_calljax.ffi.ffi_calljax.lib.xla_client.get_topology_for_devicesjax.lib.xla_client.heap_profilejax.lib.xla_client.mlir_api_versionjax.lib.xla_client.Clientjax.lib.xla_client.CompileOptionsjax.lib.xla_client.DeviceAssignmentjax.lib.xla_client.Framejax.lib.xla_client.HloShardingjax.lib.xla_client.OpShardingjax.lib.xla_client.Tracebackjax.util.HashableFunctionjax.util.as_hashable_functionjax.util.cachejax.util.safe_mapjax.util.safe_zipjax.util.split_dictjax.util.split_listjax.util.split_list_checkedjax.util.split_mergejax.util.subvalsjax.util.toposortjax.util.unzip2jax.util.wrap_namejax.util.wrapsjax.dlpack.to_dlpackjax.lax.infeedjax.lax.infeed_pjax.lax.outfeedjax.lax.outfeed_p

Breaking Changes

●Issue #1

`jax.numpy.array` no longer accepts `None`; code passing None will raise a TypeError. Replace None arguments with valid arrays or guard against None.

●Issue #2

The `config.jax_data_dependent_tracing_fallback` option has been removed; remove any usage or set appropriate tracing configuration.

●Issue #3

The `config.jax_eager_pmap` option has been removed; update code to not rely on it.

●Issue #4

Calling `lower` or `trace` AOT APIs on a `jax.jit` (or `jax.pmap`) result after additional wrappers is now disallowed. Apply `jax.jit`/`jax.pmap` as the last wrapper.

●Issue #5

The `cuda12_pip` extra has been removed; install with `pip install jax[cuda12]` instead.

Migration Steps

  1. 1
    Replace any `jax.numpy.array(None, …)` calls with valid array creation or guard against None.
  2. 2
    Remove references to `config.jax_data_dependent_tracing_fallback` and `config.jax_eager_pmap` from configuration files.
  3. 3
    If you wrap a `jax.jit` (or `jax.pmap`) result with other transformations, ensure `jax.jit`/`jax.pmap` is applied last.
  4. 4
    Install JAX with the new extra syntax, e.g., `pip install jax[cuda12]` or `pip install jax[cuda12-local]`.
  5. 5
    Update `jax.jit` calls to pass the function argument positionally and use keyword arguments for other parameters.
  6. 6
    Switch any usage of deprecated APIs to their recommended replacements (e.g., `jax.tree.unflatten`, `jax.ffi`, `jax.extend.mlir`).
  7. 7
    Upgrade your CUDA toolkit to at least 12.1 and CuDNN to 9.8 to meet the new minimum requirements.

Release Summary

This release removes several legacy tracing and configuration options, raises the minimum CUDA/CuDNN versions, updates package extras syntax, and deprecates many old APIs while introducing stricter `jax.jit` calling conventions.

Need More Details?

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

View Full Changelog