JAX v0.6.0

• Breaking changes

  • jax.numpy.array no longer accepts None. This behavior was
    deprecated since November 2023 and is now removed.
  • Removed the config.jax_data_dependent_tracing_fallback config option,
    which was added temporarily in v0.4.36 to allow users to opt out of the
    new "stackless" tracing machinery.
  • Removed the config.jax_eager_pmap config option.
  • Disallow the calling of lower and trace AOT APIs on the result
    of jax.jit if there have been subsequent wrappers applied.
    Previously this worked, but silently ignored the wrappers.
    The workaround is to apply jax.jit last among the wrappers,
    and similarly for jax.pmap.
    See #27873.
  • The cuda12_pip extra for jax has been removed; use pip install jax[cuda12]
    instead.

• Changes

  • The minimum CuDNN version is v9.8.
  • JAX is now built using CUDA 12.8. All versions of CUDA 12.1 or newer remain
    supported.
  • JAX package extras are now updated to use dash instead of underscore to
    align with PEP 685. For instance, if you were previously using pip install jax[cuda12_local]
    to install JAX, run pip install jax[cuda12-local] instead.
  • jax.jit now requires fun to be passed by position, and additional
    arguments to be passed by keyword. Doing otherwise will result in a
    DeprecationWarning in v0.6.X, and an error in starting in v0.7.X.

• Deprecations

  • jax.tree_util.build_tree is deprecated. Use jax.tree.unflatten
    instead.
  • Implemented host callback handlers for CPU and GPU devices using XLA's FFI
    and removed existing CPU/GPU handlers using XLA's custom call.
  • All APIs in jax.lib.xla_extension are now deprecated.
  • jax.interpreters.mlir.hlo and jax.interpreters.mlir.func_dialect,
    which were accidental exports, have been removed. If needed, they are
    available from jax.extend.mlir.
  • jax.interpreters.mlir.custom_call is deprecated. The APIs provided by
    jax.ffi should be used instead.
  • The deprecated use of jax.ffi.ffi_call with inline arguments is no
    longer supported. jax.ffi.ffi_call now unconditionally returns a
    callable.
  • The following exports in jax.lib.xla_client are deprecated:
    get_topology_for_devices, heap_profile, mlir_api_version, Client,
    CompileOptions, DeviceAssignment, Frame, HloSharding, OpSharding,
    Traceback.
  • The following internal APIs in jax.util are deprecated:
    HashableFunction, as_hashable_function, cache, safe_map, safe_zip,
    split_dict, split_list, split_list_checked, split_merge, subvals,
    toposort, unzip2, wrap_name, and wraps.
  • jax.dlpack.to_dlpack has been deprecated. You can usually pass a JAX
    Array directly to the from_dlpack function of another framework. If you
    need the functionality of to_dlpack, use the __dlpack__ attribute of an
    array.
  • jax.lax.infeed, jax.lax.infeed_p, jax.lax.outfeed, and
    jax.lax.outfeed_p are deprecated and will be removed in JAX v0.7.0.
  • Several previously-deprecated APIs have been removed, including:
    • From jax.lib.xla_client: ArrayImpl, FftType, PaddingType,
      PrimitiveType, XlaBuilder, dtype_to_etype,
      ops, register_custom_call_target, shape_from_pyval, Shape,
      XlaComputation.
    • From jax.lib.xla_extension: ArrayImpl, XlaRuntimeError.
    • From jax: jax.treedef_is_leaf, jax.tree_flatten, jax.tree_map,
      jax.tree_leaves, jax.tree_structure, jax.tree_transpose, and
      jax.tree_unflatten. Replacements can be found in jax.tree or
      jax.tree_util.
    • From jax.core: AxisSize, ClosedJaxpr, EvalTrace, InDBIdx, InputType,
      Jaxpr, JaxprEqn, Literal, MapPrimitive, OpaqueTraceState, OutDBIdx,
      Primitive, Token, TRACER_LEAK_DEBUGGER_WARNING, Var, concrete_aval,
      dedup_referents, escaped_tracer_error, extend_axis_env_nd, full_lower, get_referent, jaxpr_as_fun, join_effects, lattice_join,
      leaked_tracer_error, maybe_find_leaked_tracers, raise_to_shaped,
      raise_to_shaped_mappings, reset_trace_state, str_eqn_compact,
      substitute_vars_in_output_ty, typecompat, and used_axis_names_jaxpr. Most
      have no public replacement, though a few are available at jax.extend.core.
    • The vectorized argument to jax.pure_callback and
      jax.ffi.ffi_call. Use the vmap_method parameter instead.