Breaking Changes
- jax.numpy.isscalar now returns True for any array-like object with zero dimensions. Previously it only returned True for zero-dimensional array-like objects with a weak dtype.
- jax.experimental.host_callback has been deprecated since March 2024, with JAX version 0.4.26. Now we removed it. See #20385 for a discussion of alternatives.
Changes:
- jax.lax.FftType was introduced as a public name for the enum of FFT operations. The semi-public API jax.lib.xla_client.FftType has been deprecated.
- TPU: JAX now installs TPU support from the libtpu package rather than libtpu-nightly. For the next few releases JAX will pin an empty version of libtpu-nightly as well as libtpu to ease the transition; that dependency will be removed in Q1 2025.
Deprecations:
- The semi-public API jax.lib.xla_client.PaddingType has been deprecated. No JAX APIs consume this type, so there is no replacement.
- The default behavior of jax.pure_callback and jax.extend.ffi.ffi_call under vmap has been deprecated and so has the vectorized parameter to those functions. The vmap_method parameter should be used instead for better defined behavior. See the discussion in #23881 for more details.
- The semi-public API jax.lib.xla_client.register_custom_call_target has been deprecated. Use the JAX FFI instead.
- The semi-public APIs jax.lib.xla_client.dtype_to_etype, jax.lib.xla_client.ops, jax.lib.xla_client.shape_from_pyval, jax.lib.xla_client.PrimitiveType, jax.lib.xla_client.Shape, jax.lib.xla_client.XlaBuilder, and jax.lib.xla_client.XlaComputation have been deprecated. Use StableHLO instead.