Skip to content

Upgrading jax 0.7 to 0.8

Breaking0.7.2 → 0.8.3Breaking1Deprecated3New1Changed3Fixed1

Full comparisonBack to jax
BreakingUpgrade verdict

1 breaking change.

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

jax publishes an official migration guide for version 0.8. It covers more than the release notes do.

Migration guide ↗

1 package of its own release moves with it.

What you must change

Breaking1
Breaking changes.
jax.pmapjax.jitjax.shard_mapauto=jax.experimental.shard_map.shard_map__jax_array__
0.8.0GitHub ReleaseOriginalView source

Deprecations to plan for

Deprecated3
Deprecations.
jax.lax.pvaryjax.lax.pcast(..., to='varying')jax.numpy.arangejax.corecall_implget_aval
0.8.2GitHub ReleaseOriginalView source
Deprecations.
jax.sharding.PmapShardingjax.NamedShardingjx.device_put_replicatedjax.device_putjax.device_put_shardedaxis_types
0.8.1GitHub ReleaseOriginalView source
Deprecations.
jax.experimental.enable_x64jax.experimental.disable_x64jax.enable_x64jax.experimental.shard_map.shard_mapjax.shard_mapjax.experimental.pjit.pjit
0.8.0GitHub ReleaseOriginalView source