JAX v0.8.0

• Breaking changes:

  • JAX is changing the default jax.pmap implementation to one implemented in
    terms of jax.jit and jax.shard_map. jax.pmap is in maintenance mode
    and we encourage all new code to use jax.shard_map directly. See the
    migration guide for
    more information.
  • The auto= parameter of jax.experimental.shard_map.shard_map has been
    removed. This means that jax.experimental.shard_map.shard_map no longer
    supports nesting. If you want to nest shard_map calls, please use
    jax.shard_map.
  • JAX no longer allows passing objects that support __jax_array__ directly
    to, e.g. jit-ed functions. Call jax.numpy.asarray on them first.
  • jax.numpy.cov is now returns NaN for empty arrays ({jax-issue}#32305),
    and matches NumPy 2.2 behavior for single-row design matrices ({jax-issue}#32308).
  • JAX no longer accepts Array values where a dtype value is expected. Call
    .dtype on these values first.
  • The deprecated function jax.interpreters.mlir.custom_call was
    removed.
  • The jax.util, jax.extend.ffi, and jax.experimental.host_callback
    modules have been removed. All public APIs within these modules were
    deprecated and removed in v0.7.0 or earlier.
  • The deprecated symbol jax.custom_derivatives.custom_jvp_call_jaxpr_p
    was removed.
  • jax.experimental.multihost_utils.process_allgather raises an error when
    the input is a jax.Array and not fully-addressable and tiled=False. To fix
    this, pass tiled=True to your process_allgather invocation.
  • from jax.experimental.compilation_cache, the deprecated symbols
    is_initialized and initialize_cache were removed.
  • The deprecated function jax.interpreters.xla.canonicalize_dtype
    was removed.
  • jaxlib.hlo_helpers has been removed. Use jax.ffi instead.
  • The option jax_cpu_enable_gloo_collectives has been removed. Use
    jax_cpu_collectives_implementation instead.
  • The previously-deprecated interpolation argument to
    jax.numpy.percentile and jax.numpy.quantile has been
    removed; use method instead.
  • The JAX-internal for_loop primitive was removed. Its functionality,
    reading from and writing to refs in the loop body, is now directly
    supported by jax.lax.fori_loop. If you need help updating your
    code, please file a bug.
  • jax.numpy.trimzeros now errors for non-1D input.
  • The where argument to jax.numpy.sum and other reductions is now
    required to be boolean. Non-boolean values have resulted in a
    DeprecationWarning since JAX v0.5.0.
  • The deprecated functions in jax.dlpack, jax.errors,
    jax.lib.xla_bridge, jax.lib.xla_client, and
    jax.lib.xla_extension were removed.
  • jax.interpreters.mlir.dense_bool_array was removed. Use MLIR APIs to
    construct attributes instead.

• Changes

  • jax.numpy.linalg.eig now returns a namedtuple (with attributes
    eigenvalues and eigenvectors) instead of a plain tuple.
  • jax.grad and jax.vjp will now round always primals to
    float32 if float64 mode is not enabled.
  • jax.dlpack.from_dlpack now accepts arrays with non-default layouts,
    for example, transposed.
  • The default nonsymmetric eigendecomposition on NVIDIA GPUs now uses
    cusolver. The magma and LAPACK implementations are still available via the
    new implementation argument to jax.lax.linalg.eig
    ({jax-issue}#27265). The use_magma argument is now deprecated in favor
    of implementation.
  • jax.numpy.trim_zeros now follows NumPy 2.2 in supporting
    multi-dimensional inputs.

• Deprecations

  • jax.experimental.enable_x64 and jax.experimental.disable_x64
    are deprecated in favor of the new non-experimental context manager
    jax.enable_x64.
  • jax.experimental.shard_map.shard_map is deprecated; going forward use
    jax.shard_map.
  • jax.experimental.pjit.pjit is deprecated; going forward use
    jax.jit.