# JAX v0.6.0 — JAX v0.6.0

- Product: JAX (https://whatsnew.fyi/product/jax)
- Vendor: Google
- Date: 2025-04-17
- Version: v0.6.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.6.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.6.0

What's New is an index, not a publisher: every entry below links to the vendor's own release notes, which are the authoritative source. Entries are labelled where they are hand-curated sample data, pre-releases, or drawn from a secondary source such as a developer blog.

Reuse: the summaries, labels and curation here are © What's New. Quote freely with attribution and a link back; wholesale republication of the corpus is not permitted — terms: https://whatsnew.fyi/terms. The vendors' own release notes remain their publishers'.

---

- **removed** — jax.numpy.array no longer accepts None
- **removed** — Removed the config.jax_data_dependent_tracing_fallback config option
- **removed** — Removed the config.jax_eager_pmap config option
- **removed** — Removed the cuda12_pip extra for jax; use pip install jax[cuda12] instead
- **changed** — The minimum CuDNN version is v9.8
- **changed** — JAX is now built using CUDA 12.8
- **changed** — JAX package extras now use dash instead of underscore to align with PEP 685
- **changed** — jax.jit now requires fun to be passed by position, and additional arguments to be passed by keyword
- **fixed** — Disallow the calling of lower and trace AOT APIs on the result of jax.jit if there have been subsequent wrappers applied
- **deprecated** — jax.tree_util.build_tree is deprecated; use jax.tree.unflatten instead
- **deprecated** — All APIs in jax.lib.xla_extension are now deprecated
- **deprecated** — jax.interpreters.mlir.custom_call is deprecated; use APIs provided by jax.ffi instead
- **deprecated** — jax.dlpack.to_dlpack has been deprecated; use the __dlpack__ attribute of an array instead
- **deprecated** — 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
- **removed** — Removed jax.treedef_is_leaf, jax.tree_flatten, jax.tree_map, jax.tree_leaves, jax.tree_structure, jax.tree_transpose, and jax.tree_unflatten from jax
- **removed** — Removed numerous deprecated APIs from jax.lib.xla_client including ArrayImpl, FftType, PaddingType, PrimitiveType, XlaBuilder, and others
- **removed** — Removed numerous deprecated APIs from jax.core including AxisSize, ClosedJaxpr, EvalTrace, Jaxpr, JaxprEqn, Literal, Primitive, Token, and others
- **removed** — Removed the vectorized argument to jax.pure_callback and jax.ffi.ffi_call; use the vmap_method parameter instead

* 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`,

_[Truncated at 4000 characters — full notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.6.0]_