jax-ml/jax jax-v0.4.35

JAX v0.4.35

on GitHub

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.