jax-ml/jax jax-v0.8.0

JAX v0.8.0

on GitHub

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.