Skip to content

Upgrading jax 0.3 to 0.4

Breaking0.3.25 → 0.4.38Breaking3Deprecated13New7Changed10Fixed11Other2

Full comparisonBack to jax
Fetching 39 releases from GitHub…17 / 39
BreakingUpgrade verdict

3 breaking changes.

In plain terms
  • Breaking Changes.
  • Breaking changes.
  • Deprecations.
  • Deprecations & removals.
Must act Top 5 of 16 that need code or config changes

Python requirement changed in 0.4.31: >=3.7 → >=3.10.

Adds jaxlib <=0.4.38,>=0.4.38, ml-dtypes >=0.4.0, numpy >=1.26.0, opt-einsum, scipy >=1.11.1.

Installs 5 packages in all (was 0).

What you must change

Breaking3
Breaking Changes.
post_process_callnew_base_maincustom_bindconfig.jax_data_dependent_tracing_fallbackjax.experimental.jax2tf.convertnative_serialization=False
0.4.36GitHub ReleaseOriginalView source
Breaking Changes.
jax.numpy.isscalarjax.experimental.host_callback
0.4.35GitHub ReleaseOriginalView source
Breaking changes.
jax_pmap_no_rank_reductionTruearray[0]array[0:1]jax_array.addressable_shardsjax_array.addressable_data(0))
0.4.34GitHub ReleaseOriginalView source

Deprecations to plan for

Deprecated13
Deprecations.
jax.corejax.extend.corejax.extendcheck_eqncheck_typecheck_valid_jaxtype
0.4.38GitHub ReleaseOriginalView source
Deprecations.
jax.lib.xla_extension.ArrayImpljax.lib.xla_client.ArrayImpljax.Arrayjax.lib.xla_extension.XlaRuntimeErrorjax.errors.JaxRuntimeError
0.4.36GitHub ReleaseOriginalView source
Deprecations.
jax.lib.xla_client.PaddingTypejax.pure_callbackjax.extend.ffi.ffi_callvmapvectorizedvmap_method
0.4.35GitHub ReleaseOriginalView source
Deprecations.
jax.numpy.trim_zerosjax.core.pp_*jax.lib.xla_client.Devicejax.Devicejax.lib.xla_client.XlaRuntimeErrorjax.errors.JaxRuntimeError
0.4.34GitHub ReleaseOriginalView source
Deprecations.
jax.corecanonicalize_shapedimension_as_valuedefinitely_equalsymbolic_equal_dimjax.experimental.jax2tf.convert
0.4.31CHANGELOG.mdOriginalView source
Deprecations.
jax.core.pp_*TypeErrorjax.experimental.exportjax.exportxy
0.4.30CHANGELOG.mdOriginalView source
Deprecations.
jax.sharding.XLACompatibleShardingjax.sharding.Shardingjax.experimental.Exported.in_shardingsjax.experimental.Exported.in_shardings_hloout_shardingsjax.core
0.4.29GitHub ReleaseOriginalView source
Deprecations & removals.
kindjax.numpy.sortjax.numpy.argsortstable=Truestable=Falseget_compute_capability
0.4.28GitHub ReleaseOriginalView source
Deprecations & Removals.
JAX_TRITON_COMPILE_VIA_XLAjax.numpy.clipaa_mina_maxx
0.4.27CHANGELOG.mdOriginalView source
Deprecations & Removals.
jax.tree_mapjax.tree.mapjax.tree_util.tree_mapjax.clear_backendsjax.clear_cachesjax.extend.backend.clear_backends
0.4.26CHANGELOG.mdOriginalView source
Deprecations & Removals.
jax.numpy.linalg.solveapi-compatibilityjax.config.configdefine_*_stateDEFINE_*jax.config
0.4.25CHANGELOG.mdOriginalView source
Deprecations & Removals.
api-compatibilityjax.coreTracerArrayConversionErrorTracerIntegerConversionErrorUnexpectedTracerErroras_hashable_function
0.4.24CHANGELOG.mdOriginalView source
Deprecations.
device_bufferdevice_buffersarr.device_bufferarr.addressable_data(0)arr.device_buffers
0.4.22CHANGELOG.mdOriginalView source