●Change8

Migrating to JAX jax-v0.5.0

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

Released: 1/17/2025

2
Breaking Changes
10
Migration Steps
18
Affected Symbols

⚠️ Check Your Code

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

jax_threefry_partitionablejax.numpy.einsumjax.numpy.linalg.solvejax.numpy.fft.fftnjax.numpy.fft.rfftnjax.numpy.fft.ifftnjax.numpy.fft.irfftnjax.ffi.register_ffi_type_idas_text (AOT lowering)jax.interpreters.xla.abstractifyjax.interpreters.xla.pytype_aval_mappingsjax.scipy.special.lpmnjax.scipy.special.lpmn_valuesjax.extend.ffijax_enable_memoriesjax.lib.xla_client.Devicejax.lib.xla_client.XlaRuntimeErrorjax.experimental.array_api

Breaking Changes

●Issue #1

Enabled `jax_threefry_partitionable` by default, changing PRNG key semantics; update code that relied on the previous default.

●Issue #2

Dropped support for Mac x86 wheels; users on that platform must switch to Mac ARM or another supported platform.

Migration Steps

  1. 1
    Review any code that relied on the previous PRNG key behavior and adjust for the default `jax_threefry_partitionable` setting.
  2. 2
    If you were using Mac x86 wheels, switch to a supported platform (e.g., Mac ARM) or build from source.
  3. 3
    Upgrade your environment to NumPy >=1.25 and SciPy >=1.11.
  4. 4
    Replace calls to `jax.numpy.linalg.solve` with the suggested workaround for batched 1‑D RHS.
  5. 5
    Update imports: use `jax.core.abstractify` and `jax.core.pytype_aval_mappings` instead of the `jax.interpreters.xla` versions.
  6. 6
    Replace deprecated `jax.scipy.special.lpmn` and `jax.scipy.special.lpmn_values` usages or remove them.
  7. 7
    Import FFI symbols from `jax.ffi` rather than `jax.extend.ffi`.
  8. 8
    Remove usage of the `jax_enable_memories` flag; the behavior is now always enabled.
  9. 9
    Replace references to `jax.lib.xla_client.Device` and `jax.lib.xla_client.XlaRuntimeError` with `jax.Device` and `jax.errors.JaxRuntimeError`.
  10. 10
    Migrate any code using `jax.experimental.array_api` to use `jax.numpy` directly.

Release Summary

JAX meso release adds multi‑dimensional FFT support, FFI state registration, and debugging info for AOT lowering, while breaking PRNG semantics, dropping Mac x86 wheels, and raising NumPy/SciPy minimum versions.

Need More Details?

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

View Full Changelog