●Change8

Migrating to JAX jax-v0.10.0

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

Released: 4/16/2026

9
Breaking Changes
9
Migration Steps
27
Affected Symbols

⚠️ Check Your Code

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

jax.image.resizeResizeMethod.CUBIC_PYTORCHjax.lax.linalg.qrjax.lax.linalg.tridiagonal_solvejax.scipy.linalg.eigh_tridiagonaljax.numpy.ndarray.byteswapPartitionSpecjax.core.ShapedArrayjax.core.ShapedArray.vmajax.lax.linalg.tridiagonal_solvejax.numpy.clipjax.numpy.hstackjax.numpy.vstackjax.numpy.dstackjax.numpy.column_stackjax.numpy.atleast_1djax.numpy.atleast_2djax.numpy.atleast_3djax.scipy.stats.rankdatajax.ShapeDtypeStructjax.sharding.ManualAxisTypejax.distributed.initializejax.scipy.fft.dctnjax.scipy.fft.idctnjax.scipy.linalg.cho_solvejax.scipy.linalg.lu_solvejax.scipy.linalg.solve_triangular

Breaking Changes

●Issue #1

PartitionSpec objects no longer report themselves to be equal to tuples. Convert tuples to `PartitionSpec` objects before testing equality.

●Issue #2

The .vma property has been removed from `jax.core.ShapedArray`. Use .manual_axis_type.varying instead.

●Issue #3

JAX CPU devices now report their names as `cpu:0`, `cpu:1`, etc. instead of `TFRT_CPU_0`, `TFRT_CPU_1`.

●Issue #4

The config state `jax_pmap_shmap_merge` has been removed. `jax.pmap` will now always use the new implementation that wraps `jax.jit(jax.shard_map)`. Please see https://docs.jax.dev/en/latest/migrate_pmap.html for more information.

●Issue #5

`jax.device_put_sharded` and `jax.device_put_replicated` have been removed from the public API and now raise an `AttributeError` when accessed. Please see https://docs.jax.dev/en/latest/migrate_pmap.html#drop-in-replacements for drop-in replacements.

●Issue #6

The C++ pmap infrastructure has been removed. The following public APIs are no longer available: `jax.sharding.PmapSharding`, and several APIs from `jaxlib.xla_extension` and `jax.interpreters.pxla` (including `PmapFunction`, `pmap`, `NoSharding`, `Chunked`, `Unstacked`, `ShardedAxis`, `Replicated`, `ShardingSpec`, `MapTracer`, `PmapExecutable`, `parallel_callable`, `shard_args`, `xla_pmap_p`, `spec_to_indices`).

●Issue #7

The deprecated keyword arguments `a`, `a_min`, and `a_max` to `jax.numpy.clip` have been removed.

●Issue #8

Functions `jax.numpy.hstack`, `jax.numpy.vstack`, `jax.numpy.dstack`, `jax.numpy.column_stack`, `jax.numpy.atleast_1d`, `jax.numpy.atleast_2d`, and `jax.numpy.atleast_3d` no longer accept non-`ArrayLike` inputs.

●Issue #9

jax.scipy.stats.rankdata now returns floating point values in all cases, following a similar change in the SciPy 1.18 release.

Migration Steps

  1. 1
    Convert tuples to `PartitionSpec` objects before testing equality with `PartitionSpec` objects.
  2. 2
    Replace usage of the removed `.vma` property on `jax.core.ShapedArray` with `.manual_axis_type.varying`.
  3. 3
    Review migration guide at https://docs.jax.dev/en/latest/migrate_pmap.html as the config state `jax_pmap_shmap_merge` is removed and `jax.pmap` behavior has changed.
  4. 4
    Find drop-in replacements for the removed `jax.device_put_sharded` and `jax.device_put_replicated` by consulting https://docs.jax.dev/en/latest/migrate_pmap.html#drop-in-replacements.
  5. 5
    Update code that relies on the removed C++ pmap infrastructure APIs.
  6. 6
    Ensure inputs to `jax.numpy.clip` are arrays, as deprecated keyword arguments `a`, `a_min`, and `a_max` are removed.
  7. 7
    Ensure inputs to `jax.numpy.hstack`, `jax.numpy.vstack`, `jax.numpy.dstack`, `jax.numpy.column_stack`, `jax.numpy.atleast_1d`, `jax.numpy.atleast_2d`, and `jax.numpy.atleast_3d` are `ArrayLike`.
  8. 8
    Replace usage of `vma` parameter/property on `jax.ShapeDtypeStruct` with `manual_axis_type: jax.sharding.ManualAxisType` and `.manual_axis_type.varying` respectively.
  9. 9
    Update internal code using deprecated `jax.core` APIs; consider moving to `jax.extend.core` where applicable.

Release Summary

This release introduces new features like PyTorch-compatible cubic resizing and improved LAPACK parallelization, alongside significant breaking changes related to the removal of the C++ pmap infrastructure and updates to array stacking functions.

Need More Details?

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

View Full Changelog