JAX v0.4.29

Changes

• We anticipate that this will be the last release of JAX and jaxlib supporting a monolithic CUDA jaxlib. Future releases will use the CUDA plugin jaxlib (e.g. pip install jax[cuda12]).
• JAX now requires ml_dtypes version 0.4.0 or newer.
• Removed backwards-compatibility support for old usage of the jax.experimental.export API. It is not possible anymore to use from jax.experimental.export import export, and instead you should use from jax.experimental import export. The removed functionality has been deprecated since 0.4.24.

Deprecations

• jax.sharding.XLACompatibleSharding is deprecated. Please use jax.sharding.Sharding.
• jax.experimental.Exported.in_shardings has been renamed as jax.experimental.Exported.in_shardings_hlo. Same for out_shardings. The old names will be removed after 3 months.
• Removed a number of previously-deprecated APIs:
  • from {mod}jax.core: non_negative_dim, DimSize, Shape
  • from {mod}jax.lax: tie_in
  • from {mod}jax.nn: normalize
  • from {mod}jax.interpreters.xla: backend_specific_translations, translations, register_translation, xla_destructure, TranslationRule, TranslationContext, XlaOp.
• The tol argument of {func}jax.numpy.linalg.matrix_rank is being deprecated and will soon be removed. Use rtol instead.
• The rcond argument of {func}jax.numpy.linalg.pinv is being deprecated and will soon be removed. Use rtol instead.
• The deprecated jax.config submodule has been removed. To configure JAX use import jax and then reference the config object via jax.config.
• {mod}jax.random APIs no longer accept batched keys, where previously some did unintentionally. Going forward, we recommend explicit use of {func}jax.vmap in such cases.

New Functionality

• Added {func}jax.experimental.Exported.in_shardings_jax to construct shardings that can be used with the JAX APIs from the HloShardings that are stored in the Exported objects.