JAX v0.4.35

• 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.