Skip to content

Upgrading jax 0.6 to 0.7

Breaking0.6.2 → 0.7.2Breaking2Deprecated3New2Changed2Fixed1

Full comparisonBack to jax
BreakingUpgrade verdict

2 breaking changes.

In plain terms
  • Breaking changes.
  • Deprecations.
Must act 5 items need code or config changes

jax publishes official migration guides for versions 0.7, 0.7. They cover more than the release notes do.

Migration guide ↗Migration guide ↗

Python requirement changed in 0.7.0: >=3.10 → >=3.11.

1 package of its own release moves with it.

What you must change

Breaking2
Breaking changes.
jax.dlpack.from_dlpack__dlpack____dlpack_device__
0.7.2GitHub ReleaseOriginalView source
Breaking changes.
jax.stages.OutInfojax.ShapeDtypeStructjax.jitfunLayout.layout
0.7.0GitHub ReleaseOriginalView source

Deprecations to plan for

Deprecated3
Deprecations.
enable_xlanative_serializationjax2tf.convertjax_pmap_no_rank_reductionFalseTrue
0.7.2GitHub ReleaseOriginalView source
Deprecations.
jax.lax.zeros_like_arrayjax.numpy.zeros_likejax.experimental.host_callbackDeprecationWarningImportErrorNotImplementedError
0.7.1GitHub ReleaseOriginalView source
Deprecations.
jax.dlpack.SUPPORTED_DTYPESjax.dlpack.is_supported_dtypejax.scipy.special.sph_harmjax.scipy.special.sph_harm_yjax.interpreters.xlaabstractify
0.7.0GitHub ReleaseOriginalView source