# JAX changelog

> Composable transformations of NumPy programs — autodiff, JIT and vectorization.

- Vendor: Google
- Category: AI
- Official site: https://docs.jax.dev
- Tracked by: What's New (https://whatsnew.fyi/product/jax)
- Harvested from: GitHub (jax-ml/jax)
- Entries below: 25 (newest first)

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

## Releases

### v0.11.0 — JAX v0.11.0

- Date: 2026-07-16
- Version: v0.11.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.11.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.11.0

- **added** — Add hijax API documentation and helpers in jax.experimental.hijax for deriving VJPHiPrimitive autodiff rules including linearize_from_jvp, apply_derived_linearization, vjp_fwd_from_jvp, transpose_jvp, vjp_fwd_from_lin, transpose_linearized, and jvp_from_lin
- **added** — Add jax.custom_remat to the top-level jax namespace for per-function control of rematerialization under the new jax_remat3 implementation
- **changed** — Convert jax.checkpoint_policies from a namespace object to a submodule, enabling from jax.checkpoint_policies import ... syntax
- **added** — Expose SaveOnlyTheseNames, SaveAnyNamesButThese, and SaveAndOffloadOnlyTheseNames policy classes in jax.checkpoint_policies
- **added** — Add jax.Inline enum for specifying inlining policies to jax.jit
- **removed** — Remove deprecated jax.cloud_tpu_init module
- **removed** — Drop support for Python 3.11
- **removed** — Drop support for NumPy 2.0
- **removed** — Drop support for SciPy 1.14
- **removed** — Drop support for Python 3.13 free-threaded build
- **changed** — Change jax.numpy.empty and jax.numpy.empty_like to produce uninitialized arrays instead of zero-initialized arrays
- **deprecated** — Deprecate passing 2-dimensional arrays or mixed 2D and 3D arrays to jax.numpy.cross, to be removed in JAX 0.12.0
- **removed** — Remove previously-deprecated APIs from jax.core including CallPrimitive, DebugInfo, DropVar, Effect, Effects, InconclusiveDimensionOperation, JaxprTypeError, abstract_token, check_jaxpr, concrete_or_error, find_top_trace, gensym, get_opaque_trace_state, is_concrete, is_constant_dim, is_constant_shape, jaxprs_in_params, new_jaxpr_eqn, no_effects, nonempty_axis_env_DO_NOT_USE, primal_dtype_to_tangent_dtype, unsafe_am_i_under_a_jit_DO_NOT_USE, unsafe_am_i_under_a_vmap_DO_NOT_USE, unsafe_get_axis_names_DO_NOT_USE, valid_jaxtype, JaxprPpContext, JaxprPpSettings, OutputType, aval_mapping_handlers, call, concretization_function_error, custom_typechecks, literalable_types, no_axis_name, and trace_ctx
- **removed** — Remove previously-deprecated APIs from jax.interpreters.pxla including Index, MeshAxisName, MeshExecutable, global_aval_to_result_handler, global_result_handlers, are_hlo_shardings_equal, is_hlo_sharding_replicated, ArrayMapping, _UNSPECIFIED, array_mapping_to_axis_resources, and op_sharding_to_indices

* New features
  * Added a doc on defining custom derivative rules with the experimental
    hijax API (`hijax-custom-derivatives`), along with
    `jax.experimental.hijax` helpers for deriving `VJPHiPrimitive` autodiff
    rules from a `jvp` or `lin` rule: `linearize_from_jvp` with
    `apply_derived_linearization`, `vjp_fwd_from_jvp` with `transpose_jvp`,
    `vjp_fwd_from_lin` with `transpose_linearized`, and `jvp_from_lin`.
  * Added `jax.custom_remat` to the top-level `jax` namespace, for
    per-function control of rematerialization under the new `jax_remat3`
    implementation.
  * `jax.checkpoint_policies` is now a submodule rather than a namespace
    object (so `from jax.checkpoint_policies import ...` now works; attribute
    access is unchanged), and it additionally exposes the name-based policy
    classes `SaveOnlyTheseNames`, `SaveAnyNamesButThese`, and
    `SaveAndOffloadOnlyTheseNames`.
  * Added `jax.Inline` enum for specify inlining policies to
    `jax.jit`.

* Breaking changes
  * The deprecated module j`ax.cloud_tpu_init` was removed. This did nothing and
    references to it can be safely removed.
  * Support for Python 3.11, NumPy 2.0, and SciPy 1.14 has been dropped, per the
    [deprecation policy](https://docs.jax.dev/en/latest/deprecation.html).
  * Support for Python 3.13 free-threaded (`3.13t`) has been dropped. Python
    3.13 free-threaded was an experimental build needed to bootstrap free
    threading support. Now that Python 3.14t is stable, it is time to drop the
    experimental build. Other parts of the Python ecosystem (e.g.
    `cibuildwheel`, `scipy`) are making similar moves.
  * `jax.numpy.empty` and `jax.numpy.empty_like` now produce
    uninitialized arrays, similar to their NumPy counterparts. Prior to v0.11.0,
    they produced arrays initialized to zeros. To recover the previous behavior,
    use `jax.numpy.zeros` or `jax.numpy.zeros_like` instead.

* Deprecations
  * Passing 2-dimensional arrays (or mixed 2D and 3D arrays) to `jax.numpy.cross` is deprecated and will be removed in JAX 0.12.0, aligning with NumPy 2.5 behavior.
  * Several previously-deprecated APIs from `jax.core` have been removed, including
    `CallPrimitive`, `DebugInfo`, `DropVar`, `Effect`, `Effects`, `InconclusiveDimensionOperation`,
    `JaxprTypeError`, `abstract_token`, `check_jaxpr`, `concrete_or_error`, `find_top_trace`, `gensym`,
    `get_opaque_trace_state`, `is_concrete`, `is_constant_dim`, `is_constant_shape`, `jaxprs_in_params`,
    `new_jaxpr_eqn`, `no_effects`, `nonempty_axis_env_DO_NOT_USE`, `primal_dtype_to_tangent_dtype`,
    `unsafe_am_i_under_a_jit_DO_NOT_USE`, `unsafe_am_i_under_a_vmap_DO_NOT_USE`,
    `unsafe_get_axis_names_DO_NOT_USE`, `valid_jaxtype`, `JaxprPpContext`, `JaxprPpSettings`, `OutputType`,
    `aval_mapping_handlers`, `call`, `concretization_function_error`, `custom_typechecks`,
    `literalable_types`, `no_axis_name`, and `trace_ctx`.
  * Several previously-deprecated APIs from `jax.interpreters.pxla` have been removed, including
    `Index`, `MeshAxisName`, `MeshExecutable`, `global_aval_to_result_handler`, `global_result_handlers`,
    `are_hlo_shardings_equal`, `is_hlo_sharding_replicated`, `ArrayMapping`, `_UNSPECIFIED`,
    `array_mapping_to_axis_resources`, and `op_sharding_to_indices`.

### v0.10.2 — JAX v0.10.2

- Date: 2026-06-17
- Version: v0.10.2
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.10.2
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.10.2

- **added** — Added jax.scipy.linalg.invhilbert for the closed-form inverse of the Hilbert matrix
- **added** — Added jax.scipy.linalg.invpascal for the inverse of the Pascal matrix
- **added** — Added jax.scipy.linalg.fiedler_companion for constructing the pentadiagonal Fiedler companion matrix of a polynomial
- **added** — Added jax.ShapeDtypeStruct.like as a shortcut for constructing a jax.ShapeDtypeStruct from an object with shape and dtype attributes
- **changed** — Improved JAX tracing performance

* New features
  * Added `jax.scipy.linalg.invhilbert` for the closed-form inverse
    of the Hilbert matrix ({jax-issue}`#10144`).
  * Added `jax.scipy.linalg.invpascal` for the inverse of the Pascal
    matrix (`#10144`).
  * Added `jax.scipy.linalg.fiedler_companion` for constructing the
    pentadiagonal Fiedler companion matrix of a polynomial
    (`#10144`).
  * Added `jax.ShapeDtypeStruct.like` -- a shortcut for constructing a
    `jax.ShapeDtypeStruct` from an object with `shape` and `dtype`
    attributes.
* We have made a number of improvements to JAX tracing performance.

### v0.10.1 — JAX v0.10.1

- Date: 2026-05-20
- Version: v0.10.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.10.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.10.1

- **added** — Added ResizeMethod.AREA to jax.image.resize, which matches TensorFlow's AREA resizing
- **added** — Added jax.scipy.linalg.hadamard for constructing Hadamard matrices
- **added** — Added jax.scipy.linalg.circulant for constructing circulant matrices
- **added** — Added jax.scipy.linalg.dft for constructing discrete Fourier transform matrices
- **added** — Added jax.scipy.linalg.leslie for constructing Leslie matrices
- **added** — Added jax.scipy.linalg.companion for constructing companion matrices from polynomial coefficients
- **added** — Added jax.scipy.linalg.fiedler for constructing symmetric Fiedler matrices
- **added** — Added jax.scipy.linalg.helmert for constructing Helmert matrices
- **added** — Added jax.random.key_dtype to get the dtype corresponding to a PRNG implementation name
- **changed** — jax.random.key and wrap_key_data now accept a dtype argument
- **deprecated** — with mesh: context manager has been deprecated, use with jax.set_mesh(mesh): instead
- **deprecated** — Passing the copy, order, and ndmin arguments to jax.numpy.array positionally is deprecated, use keyword arguments instead
- **deprecated** — Python dict_values, generators, zip return type and iterators are deprecated by default when used as leaves in pytrees

* New features
  * Added `ResizeMethod.AREA` to `jax.image.resize`, which matches
    TensorFlow's AREA resizing (#20098).
  * Added `jax.scipy.linalg.hadamard` for constructing Hadamard
    matrices (#10144).
  * Added `jax.scipy.linalg.circulant` for constructing circulant
    matrices (#10144).
  * Added `jax.scipy.linalg.dft` for constructing discrete Fourier
    transform matrices (#10144).
  * Added `jax.scipy.linalg.leslie` for constructing Leslie matrices
    (#10144).
  * Added `jax.scipy.linalg.companion` for constructing companion
    matrices from polynomial coefficients (#10144).
  * Added `jax.scipy.linalg.fiedler` for constructing symmetric Fiedler
    matrices (#10144).
  * Added `jax.scipy.linalg.helmert` for constructing Helmert matrices
    (#10144).
  * Moved RNG APIs from "implementations" to dtypes (#27854):
    * Added `jax.random.key_dtype` to get the dtype corresponding to a PRNG
      implementation name.
    * `jax.random.key` and `wrap_key_data` now accept a `dtype` argument.

* Breaking changes
  * `with mesh:` context manager has been deprecated. Please use
    `with jax.set_mesh(mesh):` instead.

* Deprecations
  * Passing the `copy`, `order`, and `ndmin` arguments to
    `jax.numpy.array` positionally is deprecated. Use keyword arguments
    instead. This matches the signature of `numpy.array`.
  * Python `dict_values`, generators, zip return type and iterators generally
    are deprecated by default when used as leaves in pytrees. In a future
    version of JAX, this will become an error, if you depend on using them as
    leaves, pass `is_leaf` to `jax.tree.*` methods.

### v0.10.0 — JAX v0.10.0

- Date: 2026-04-16
- Version: v0.10.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.10.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.10.0

- **added** — Add ResizeMethod.CUBIC_PYTORCH to jax.image.resize to match PyTorch's bicubic resize
- **added** — Support differentiation of jax.lax.linalg.qr for wide matrices and when full_matrices is True
- **added** — Parallelize LAPACK operations along the batch dimension on CPU
- **added** — Add perturb_singular argument to jax.lax.linalg.tridiagonal_solve to handle singular matrices by perturbing near-zero pivots in the LU decomposition
- **added** — Support computing eigenvectors on CPU and GPU in jax.scipy.linalg.eigh_tridiagonal
- **added** — Add jax.numpy.ndarray.byteswap method
- **removed** — Remove PartitionSpec equality with tuples
- **removed** — Remove .vma property from jax.core.ShapedArray in favor of .manual_axis_type.varying
- **changed** — JAX CPU devices now report their names as cpu:0, cpu:1, etc. instead of TFRT_CPU_0, TFRT_CPU_1
- **removed** — Remove config state jax_pmap_shmap_merge; jax.pmap now always uses the new implementation that wraps jax.jit(jax.shard_map)
- **removed** — Remove jax.device_put_sharded and jax.device_put_replicated from the public API
- **removed** — Remove C++ pmap infrastructure including jax.sharding.PmapSharding and related APIs from jaxlib.xla_extension and jax.interpreters.pxla
- **removed** — Remove deprecated keyword arguments a, a_min, and a_max from jax.numpy.clip
- **removed** — Remove support for non-ArrayLike inputs to jax.numpy.hstack, jax.numpy.vstack, jax.numpy.dstack, jax.numpy.column_stack, jax.numpy.atleast_1d, jax.numpy.atleast_2d, and jax.numpy.atleast_3d
- **changed** — jax.scipy.stats.rankdata now returns floating point values in all cases, following SciPy 1.18
- **changed** — Increase minimum supported SciPy version to 1.14
- **changed** — Replace vma parameter of jax.ShapeDtypeStruct with manual_axis_type: jax.sharding.ManualAxisType
- **removed** — Remove experimental jax.experimental.custom_dce.custom_dce
- **fixed** — Fix a bug that led to differing output between CPU and GPU for non-symmetric multidimensional IRFFTs
- **fixed** — Fix an error when tiny matrices were passed to jax.lax.linalg.tridiagonal_solve on GPU
- **fixed** — Fix a bug in jax.scipy.fft.dctn and idctn where axes=None incorrectly defaulted to all axes when s was specified
- **fixed** — Fix jax.distributed.initialize() on a GCE TPU Managed Instance Group raising an IndexError

* New features:
  * Added `ResizeMethod.CUBIC_PYTORCH` to jax.image.resize to match
    PyTorch's bicubic resize (#15768).
  * We now support differentiation of jax.lax.linalg.qr for wide
    matrices and when `full_matrices` is `True`.
  * LAPACK operations are now parallelized along the batch dimension on CPU.
  * Added `perturb_singular` argument to
    jax.lax.linalg.tridiagonal_solve to handle singular matrices by
    perturbing near-zero pivots in the LU decomposition. This is useful for
    solving numerically singular systems when computing eigenvectors by inverse
    iteration.
  * jax.scipy.linalg.eigh_tridiagonal now supports computing
    eigenvectors on CPU and GPU.
  * Added the jax.numpy.ndarray.byteswap method.

* Breaking changes:
  * `PartitionSpec` objects no longer report themselves to be equal to tuples.
    Convert tuples to `PartitionSpec` objects before testing equality.
  * The `.vma` property has been removed from `jax.core.ShapedArray`. Use
    `.manual_axis_type.varying` instead.
  * JAX CPU devices now report their names as `cpu:0`, `cpu:1`, etc. instead of
    `TFRT_CPU_0`, `TFRT_CPU_1`.
  * The config state `jax_pmap_shmap_merge` has been removed. `jax.pmap`
    will now always use the new implementation that wraps
    `jax.jit(jax.shard_map)`. Please see
    https://docs.jax.dev/en/latest/migrate_pmap.html for more information.
  * `jax.device_put_sharded` and `jax.device_put_replicated` have been removed
    from the public API and now raise an `AttributeError` when accessed.
    Please see
    https://docs.jax.dev/en/latest/migrate_pmap.html#drop-in-replacements for
    drop-in replacements.
  * The C++ pmap infrastructure has been removed. The following public APIs
    are no longer available:
    * `jax.sharding.PmapSharding`
    * From `jaxlib.xla_extension`: `PmapFunction`, `pmap`,
      `NoSharding`, `Chunked`, `Unstacked`, `ShardedAxis`, `Replicated`,
      `ShardingSpec`.
    * From `jax.interpreters.pxla`: `MapTracer`, `PmapExecutable`,
      `parallel_callable`, `shard_args`, `xla_pmap_p`, `Chunked`,
      `NoSharding`, `Replicated`, `ShardedAxis`, `ShardingSpec`,
      `Unstacked`, `spec_to_indices`.
  * The deprecated keyword arguments `a`, `a_min`, and `a_max` to
    `jax.numpy.clip` have been removed.
  * Functions `jax.numpy.hstack`, `jax.numpy.vstack`, `jax.numpy.dstack`,
    `jax.numpy.column_stack`, `jax.numpy.atleast_1d`, `jax.numpy.atleast_2d`,
    and `jax.numpy.atleast_3d` no longer accept non-`ArrayLike` inputs.
    Doing so previously issued a `DeprecationWarning`.
  * jax.scipy.stats.rankdata now returns floating point values in
    all cases, following a similar change in the SciPy 1.18 release.

* Deprecations:
  * A number of internal APIs in `jax.core` have been newly deprecated and
    some have been moved to `jax.extend.core`. These include `CallPrimitive`,
    `DebugInfo`, `DropVar`, `Effect`, `Effects`, `InconclusiveDimensionOperation`,
    `JaxprTypeError`, `check_jaxpr`, `concrete_or_error`, `find_top_trace`,
    `gensym`, `get_opaque_trace_state`, `jaxprs_in_params`, `new_jaxpr_eqn`,
    `no_effects`, `nonempty_axis_env_DO_NOT_USE`, `primal_dtype_to_tangent_dtype`,
    `unsafe_am_i_under_a_jit_DO_NOT_USE`, `unsafe_am_i_under_a_vmap_DO_NOT_USE`,
    `unsafe_get_axis_names_DO_NOT_USE`, `valid_jaxtype`, `JaxprPpContext`,
    `JaxprPpSettings`, `OutputType`, `abstract_token`, `aval_mapping_handlers`,
    `call`, `concretization_function_error`, `custom_typechecks`, `is_concrete`,
    `is_constant_dim`, `is_constant_shape`, `literalable_types`, `no_axis_name`,
    `pytype_aval_mappings`, and `trace_ctx`.

* Changes:
  * The minimum supported SciPy version is now 1.14.
  * `vma` parameter of `jax.ShapeDtypeStruct` has been replaced with
    `manual_axis_type: jax.sharding.ManualAxisType`. The `.vma` property has
    been replaced with `.manual_axis_type.varying`.
  * Removed experimental jax.exper

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

### v0.9.2 — JAX v0.9.2

- Date: 2026-03-18
- Version: v0.9.2
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.9.2
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.9.2

- **changed** — The semi-private type jax._src.literals.TypedNdArray is now a subclass of np.ndarray, rather than a duck type of it
- **changed** — jax.numpy.arange with step specified no longer generates the array on host, resulting in more efficient code but potentially less precise outputs for narrow-width floats like bfloat16

##### JAX 0.9.2 (March 2, 2026)

* Changes:
  * The semi-private type `jax._src.literals.TypedNdArray` is now a subclass of
    `np.ndarray`, rather than a duck type of it.
  * `jax.numpy.arange` with `step` specified no longer generates the array
    on host. The benefit is more efficient code, though this can lead to less
    precise outputs for narrow-width floats (e.g. bfloat16). To recover the
    previous behavior in this case, use `jnp.array(np.arange(...))`.

### v0.9.1 — JAX v0.9.1

- Date: 2026-03-02
- Version: v0.9.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.9.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.9.1

- **changed** — JAX tracers that are not of Array type (e.g., of Ref type) will no longer report themselves to be instances of Array
- **changed** — Using jax.shard_map in Explicit mode will raise an error if the PartitionSpec of input does not match the PartitionSpec specified in in_specs, acting like an assert instead of an implicit reshard
- **added** — Added a debug config jax_compilation_cache_check_contents that verifies cache contents match when a value is put() and causes misses when get() is called on a value not put() by the current process

* Changes:
  * JAX tracers that are not of `Array` type (e.g., of `Ref` type) will no
    longer report themselves to be instances of `Array`.
  * Using `jax.shard_map` in Explicit mode will raise an error
    if the PartitionSpec of input does not match the PartitionSpec specified in
    `in_specs`. In other words, it will act like an assert instead of an
    implicit reshard.
    `in_specs` is an optional argument so you can omit specifying it
    and `shard_map` will infer the `PartitionSpec` from the argument. If you
    want to reshard your inputs, you can use `jax.reshard` on the arguments and
    then pass those args to shard_map.

* New features:
  * Added a debug config `jax_compilation_cache_check_contents`. If set, we miss
    when `get()` is called on a value that has not been `put()` by the current
    process, even if the value is actually in the disk cache. When a value is
    `put()`, we verify that its contents match.

### v0.9.0.1 — JAX v0.9.0.1

- Date: 2026-02-05
- Version: v0.9.0.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.9.0.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.9.0.1

JAX v0.9.0.1 is identical to v0.9.0 with the commits from the following four PRs patched in:

- https://github.com/openxla/xla/pull/36579
- https://github.com/openxla/xla/pull/36345
- https://github.com/openxla/xla/pull/36755
- https://github.com/openxla/xla/pull/36696

### v0.8.3 — JAX v0.8.3

- Date: 2026-01-29
- Version: v0.8.3
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.8.3
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.8.3

- **fixed** — Fixed issue from XLA commit 4bc723da9766d784920f4e566f87afc6ffbf6a5b
- **fixed** — Fixed issue from XLA commit 21552fd83ec0f02ec3f418f4ac356bfa1d91ce4d

JAX v0.8.3 is identical to v0.8.2 with the following two bug fixes patched in:

- https://github.com/openxla/xla/commit/4bc723da9766d784920f4e566f87afc6ffbf6a5b
- https://github.com/openxla/xla/commit/21552fd83ec0f02ec3f418f4ac356bfa1d91ce4d

### v0.9.0 — JAX v0.9.0

- Date: 2026-01-20
- Version: v0.9.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.9.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.9.0

- **added** — Added jax.thread_guard, a context manager that detects when devices are used by multiple threads in multi-controller JAX
- **fixed** — Fixed a workspace size calculation error for pivoted QR (magma_zgeqp3_gpu) in MAGMA 2.9.0 when using use_magma=True and pivoting=True
- **removed** — Removed the flag jax_collectives_common_channel_id
- **removed** — Removed the jax_pmap_no_rank_reduction config state; the no-rank-reduction behavior is now the only supported behavior
- **deprecated** — Setting the jax_pmap_shmap_merge config state is deprecated and will be removed in JAX v0.10.0
- **deprecated** — jax.numpy.fix is deprecated; jax.numpy.trunc is a drop-in replacement
- **changed** — jax.export now supports explicit sharding with a new export serialization format version that includes NamedSharding, the abstract mesh, and the partition spec
- **changed** — When calling exported modules, the abstract mesh must now match the one used at export time, including the axis names

* New features:

  * Added `jax.thread_guard`, a context manager that detects when devices
    are used by multiple threads in multi-controller JAX.

* Bug fixes:
  * Fixed a workspace size calculation error for pivoted QR (`magma_zgeqp3_gpu`)
    in MAGMA 2.9.0 when using `use_magma=True` and `pivoting=True`.
    (#34145).

* Deprecations:
  * The flag `jax_collectives_common_channel_id` was removed.
  * The `jax_pmap_no_rank_reduction` config state has been removed. The
    no-rank-reduction behavior is now the only supported behavior: a
    `jax.pmap`ped function `f` sees inputs of the same rank as the input to
    `jax.pmap(f)`. For example, if `jax.pmap(f)` receives shape `(8, 128)` on
    8 devices, then `f` receives shape `(1, 128)`.
  * Setting the `jax_pmap_shmap_merge` config state is deprecated in JAX v0.9.0
    and will be removed in JAX v0.10.0.
  * `jax.numpy.fix` is deprecated, anticipating the deprecation of
    `numpy.fix` in NumPy v2.5.0. `jax.numpy.trunc` is a drop-in
    replacement.

* Changes:
  * `jax.export` now supports explicit sharding. This required a new
    export serialization format version that includes the NamedSharding,
    including the abstract mesh, and the partition spec. As part of this
    change we have added a restriction in the use of exported modules: when
    calling them the abstract mesh must match the one used at export time,
    including the axis names. Previously, only the number of the devices
    mattered.

### v0.8.2 — JAX v0.8.2

- Date: 2025-12-18
- Version: v0.8.2
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.8.2
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.8.2

- **deprecated** — jax.lax.pvary has been deprecated; use jax.lax.pcast(..., to='varying') as the replacement
- **deprecated** — Complex arguments passed to jax.numpy.arange now result in a deprecation warning
- **deprecated** — From jax.core, the following symbols are newly deprecated: call_impl, get_aval, mapped_aval, subjaxprs, set_current_trace, take_current_trace, traverse_jaxpr_params, unmapped_aval, AbstractToken, and TraceTag
- **deprecated** — All symbols in jax.interpreters.pxla are deprecated
- **changed** — jax's Tracer no longer inherits from jax.Array at runtime, but jax.Array now uses a custom metaclass such that isinstance(x, Array) is true if an object x represents a traced Array
- **removed** — jax.experimental.si_vjp has been deleted; use jax.vjp instead

* Deprecations
  * `jax.lax.pvary` has been deprecated.
    Please use `jax.lax.pcast(..., to='varying')` as the replacement.
  * Complex arguments passed to `jax.numpy.arange` now result in a
    deprecation warning, because the output is poorly-defined.
  * From `jax.core` a number of symbols are newly deprecated including:
    `call_impl`, `get_aval`, `mapped_aval`, `subjaxprs`, `set_current_trace`,
    `take_current_trace`, `traverse_jaxpr_params`, `unmapped_aval`,
    `AbstractToken`,  and `TraceTag`.
  * All symbols in `jax.interpreters.pxla` are deprecated. These are
    primarily JAX internal APIs, and users should not rely on them.

* Changes:
  * jax's `Tracer` no longer inherits from `jax.Array` at runtime. However,
    `jax.Array` now uses a custom metaclass such `isinstance(x, Array)` is true
    if an object `x` represents a traced `Array`. Only some `Tracer`s represent
    `Array`s, so it is not correct for `Tracer` to inherit from `Array`.

    For the moment, during Python type checking, we continue to declare `Tracer`
    as a subclass of `Array`, however we expect to remove this in a future
    release.
  * `jax.experimental.si_vjp` has been deleted.
    `jax.vjp` subsumes it's functionality.

### v0.8.1 — JAX v0.8.1

- Date: 2025-11-18
- Version: v0.8.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.8.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.8.1

- **added** — jax.jit now supports the decorator factory pattern, allowing @jax.jit(static_argnames=['n']) syntax instead of requiring functools.partial
- **changed** — jax.lax.linalg.eigh now accepts an implementation argument to select between QR (CPU/GPU), Jacobi (GPU/TPU), and QDWH (TPU) implementations
- **changed** — EighImplementation enum is now publicly exported from jax.lax.linalg
- **changed** — jax.lax.linalg.svd now implements an algorithm that uses polar decomposition on CUDA GPUs
- **fixed** — Fixed a bug introduced in JAX 0.7.2 where eigh failed for large matrices on GPU
- **deprecated** — jax.sharding.PmapSharding is deprecated; use jax.NamedSharding instead
- **deprecated** — jax.device_put_replicated is deprecated; use jax.device_put with the appropriate sharding instead
- **deprecated** — jax.device_put_sharded is deprecated; use jax.device_put with the appropriate sharding instead
- **deprecated** — Default axis_types of jax.make_mesh will change in JAX v0.9.0 to return jax.sharding.AxisType.Explicit; leaving axis_types unspecified will raise a DeprecationWarning
- **deprecated** — jax.cloud_tpu_init and its contents are deprecated

* New features:

  * `jax.jit` now supports the decorator factory pattern; i.e instead of
    writing
    ```python
    @functools.partial(jax.jit, static_argnames=['n'])
    def f(x, n):
      ...
    ```
    you may write
    ```python
    @jax.jit(static_argnames=['n'])
    def f(x, n):
      ...
    ```

* Changes:

  * `jax.lax.linalg.eigh` now accepts an `implementation` argument to
    select between QR (CPU/GPU), Jacobi (GPU/TPU), and QDWH (TPU)
    implementations. The `EighImplementation` enum is publicly exported from
    `jax.lax.linalg`.

  * `jax.lax.linalg.svd` now implements an `algorithm` that uses the polar
    decomposition on CUDA GPUs. This is also an alias for the existing algorithm
    on TPUs.

* Bug fixes:

  * Fixed a bug introduced in JAX 0.7.2 where eigh failed for large matrices on
    GPU (#33062).

* Deprecations:
  * `jax.sharding.PmapSharding` is now deprecated. Please use
    `jax.NamedSharding` instead.
  * `jx.device_put_replicated` is now deprecated. Please use `jax.device_put`
    with the appropriate sharding instead.
  * `jax.device_put_sharded` is now deprecated. Please use `jax.device_put` with
    the appropriate sharding instead.
  * Default `axis_types` of `jax.make_mesh` will change in JAX v0.9.0 to return
  `jax.sharding.AxisType.Explicit`. Leaving axis_types unspecified will raise a
  `DeprecationWarning`.
  * `jax.cloud_tpu_init` and its contents were deprecated. There is no reason for a user to import or use the contents of this module; JAX handles this for you automatically if needed.

### v0.8.0 — JAX v0.8.0

- Date: 2025-10-15
- Version: v0.8.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.8.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.8.0

- **changed** — JAX is changing the default jax.pmap implementation to one implemented in terms of jax.jit and jax.shard_map
- **removed** — The auto= parameter of jax.experimental.shard_map.shard_map has been removed
- **removed** — JAX no longer allows passing objects that support __jax_array__ directly to jit-ed functions
- **changed** — jax.numpy.cov now returns NaN for empty arrays and matches NumPy 2.2 behavior for single-row design matrices
- **removed** — JAX no longer accepts Array values where a dtype value is expected
- **removed** — The deprecated function jax.interpreters.mlir.custom_call was removed
- **removed** — The jax.util, jax.extend.ffi, and jax.experimental.host_callback modules have been removed
- **removed** — The deprecated symbol jax.custom_derivatives.custom_jvp_call_jaxpr_p was removed
- **changed** — jax.experimental.multihost_utils.process_allgather raises an error when the input is a jax.Array and not fully-addressable and tiled=False
- **removed** — From jax.experimental.compilation_cache, the deprecated symbols is_initialized and initialize_cache were removed
- **removed** — The deprecated function jax.interpreters.xla.canonicalize_dtype was removed
- **removed** — jaxlib.hlo_helpers has been removed
- **removed** — The option jax_cpu_enable_gloo_collectives has been removed
- **removed** — The previously-deprecated interpolation argument to jax.numpy.percentile and jax.numpy.quantile has been removed
- **removed** — The JAX-internal for_loop primitive was removed
- **changed** — jax.numpy.trimzeros now errors for non-1D input
- **changed** — The where argument to jax.numpy.sum and other reductions is now required to be boolean
- **removed** — The deprecated functions in jax.dlpack, jax.errors, jax.lib.xla_bridge, jax.lib.xla_client, and jax.lib.xla_extension were removed
- **removed** — jax.interpreters.mlir.dense_bool_array was removed
- **changed** — jax.numpy.linalg.eig now returns a namedtuple with attributes eigenvalues and eigenvectors instead of a plain tuple
- **changed** — jax.grad and jax.vjp will now round primals to float32 if float64 mode is not enabled

* 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](https://docs.jax.dev/en/latest/migrate_pmap.html) 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 depreca

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

### v0.7.2 — JAX v0.7.2

- Date: 2025-09-16
- Version: v0.7.2
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.7.2
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.7.2

- **removed** — jax.dlpack.from_dlpack no longer accepts a DLPack capsule and must be called with an array implementing __dlpack__ and __dlpack_device__
- **changed** — Minimum supported NumPy version is now 2.0
- **changed** — Minimum supported SciPy version is now 1.13
- **changed** — JAX now represents constants in its internal jaxpr representation as a LiteralArray, a private JAX type that duck types as a numpy.ndarray
- **fixed** — arr.view(dtype=None) now returns the array unchanged, matching NumPy's semantics
- **fixed** — jax.random.randint now produces a less-biased distribution for 8-bit and 16-bit integer types
- **deprecated** — The parameters enable_xla and native_serialization for jax2tf.convert are deprecated and will be removed in a future version
- **deprecated** — Setting the config state jax_pmap_no_rank_reduction to False is deprecated

* Breaking changes:

  * `jax.dlpack.from_dlpack` no longer accepts a DLPack capsule. This
    behavior was deprecated and is now removed. The function must be called
    with an array implementing `__dlpack__` and `__dlpack_device__`.

* Changes
  * The minimum supported NumPy version is now 2.0. Since SciPy 1.13 is required
    for NumPy 2.0 support, the minimum supported SciPy version is now 1.13.

  * JAX now represents constants in its internal jaxpr representation as a
    `LiteralArray`, which is a private JAX type that duck types as a
    `numpy.ndarray`. This type may be exposed to users via `custom_jvp` rules,
    for example, and may break code that uses `isinstance(x, np.ndarray)`. If
    this breaks your code, you may convert these arrays to classic NumPy arrays
    using `np.asarray(x)`.

* Bug fixes
  * `arr.view(dtype=None)` now returns the array unchanged, matching NumPy's
    semantics. Previously it returned the array with a float dtype.
  * `jax.random.randint` now produces a less-biased distribution for 8-bit and
    16-bit integer types ({jax-issue}`#27742`). To restore the previous biased
    behavior, you may temporarily set the `jax_safer_randint` configuration to
    `False`, but note this is a temporary config that will be removed in a
    future release.

* Deprecations:
  * The parameters `enable_xla` and `native_serialization` for `jax2tf.convert`
    are deprecated and will be removed in a future version of JAX. These were
    used for jax2tf with non-native serialization, which has been now removed.
  * Setting the config state `jax_pmap_no_rank_reduction` to `False` is
    deprecated. By default, `jax_pmap_no_rank_reduction` will be set to `True`
    and `jax.pmap` shards will not have their rank reduced, keeping the same
    rank as their enclosing array.

### v0.7.1 — JAX v0.7.1

- Date: 2025-08-20
- Version: v0.7.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.7.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.7.1

- **added** — JAX now ships Python 3.14 and 3.14t wheels
- **added** — JAX now ships Python 3.13t and 3.14t wheels on Mac
- **changed** — Exposed jax.set_mesh which acts as a global setter and a context manager
- **removed** — Removed jax.sharding.use_mesh in favor of jax.set_mesh
- **changed** — JAX is now built using CUDA 12.9
- **changed** — jax.lax.dot now implements the general dot product via the optional dimension_numbers argument
- **deprecated** — jax.lax.zeros_like_array is deprecated in favor of jax.numpy.zeros_like
- **deprecated** — Attempting to import jax.experimental.host_callback now results in a DeprecationWarning and will result in an ImportError starting in JAX v0.8.0
- **deprecated** — In jax.lax.dot, passing the precision and preferred_element_type arguments by position is deprecated; pass them by explicit keyword instead
- **deprecated** — Several internal APIs have been deprecated from jax.interpreters.ad, jax.interpreters.batching, and jax.interpreters.partial_eval

* New features
  * JAX now ships Python 3.14 and 3.14t wheels.
  * JAX now ships Python 3.13t and 3.14t wheels on Mac. Previously we only
    offered free-threading builds on Linux.

* Changes
  * Exposed `jax.set_mesh` which acts as a global setter and a context manager.
    Removed `jax.sharding.use_mesh` in favor of `jax.set_mesh`.
  * JAX is now built using CUDA 12.9. All versions of CUDA 12.1 or newer remain
    supported.
  * `jax.lax.dot` now implements the general dot product via the optional
    ``dimension_numbers`` argument.

* Deprecations:

  * `jax.lax.zeros_like_array` is deprecated. Please use
    `jax.numpy.zeros_like` instead.
  * Attempting to import `jax.experimental.host_callback` now results in
    a `DeprecationWarning`, and will result in an `ImportError` starting in JAX
    v0.8.0. Its APIs have raised `NotImplementedError` since JAX version 0.4.35.
  * In `jax.lax.dot`, passing the ``precision`` and ``preferred_element_type``
    arguments by position is deprecated. Pass them by explicit keyword instead.
  * Several dozen internal APIs have been deprecated from `jax.interpreters.ad`,
    `jax.interpreters.batching`, and `jax.interpreters.partial_eval`; they
    are used rarely if ever outside JAX itself, and most are deprecated without any
    public replacement.

### v0.7.0 — JAX v0.7.0

- Date: 2025-07-22
- Version: v0.7.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.7.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.7.0

- **added** — Added jax.P which is an alias for jax.sharding.PartitionSpec
- **added** — Added jax.tree.reduce_associative
- **changed** — JAX is migrating from GSPMD to Shardy by default
- **changed** — JAX autodiff is switching to using direct linearization by default instead of implementing linearization via JVP and partial eval
- **removed** — jax.stages.OutInfo has been replaced with jax.ShapeDtypeStruct
- **changed** — jax.jit now requires fun to be passed by position and additional arguments to be passed by keyword
- **changed** — The minimum Python version is now 3.11
- **changed** — Layout, .layout, .input_layouts and .output_layouts have been renamed to Format, .format, .input_formats and .output_formats
- **changed** — DeviceLocalLayout and .device_local_layout have been renamed to Layout and .layout
- **removed** — jax.experimental.shard module has been deleted and all APIs have been moved to jax.sharding endpoint
- **removed** — lax.infeed and lax.outfeed were removed after being deprecated in JAX 0.6
- **removed** — transfer_to_infeed and transfer_from_outfeed methods were removed from Device objects
- **changed** — jax.extend.core.primitives.pjit_p primitive has been renamed to jit_p and its name attribute has changed from pjit to jit
- **removed** — jax.extend.backend.add_clear_backends_callback has been removed
- **deprecated** — jax.dlpack.SUPPORTED_DTYPES is deprecated; use jax.dlpack.is_supported_dtype instead
- **deprecated** — jax.scipy.special.sph_harm has been deprecated; use jax.scipy.special.sph_harm_y instead
- **deprecated** — jax.interpreters.xla.canonicalize_dtype is deprecated; prefer jax.dtypes.canonicalize_dtype for canonicalizing dtypes and jax.core.valid_jaxtype for checking valid jax input
- **deprecated** — jax.lib.xla_bridge.get_compile_options is deprecated and replaced by jax.extend.backend.get_compile_options
- **removed** — jax.extend.ffi was removed after being deprecated in v0.5.0; use jax.ffi instead

* New features:
  * Added `jax.P` which is an alias for `jax.sharding.PartitionSpec`.
  * Added `jax.tree.reduce_associative`.

* Breaking changes:
  * JAX is migrating from GSPMD to Shardy by default. See the
    [migration guide](https://docs.jax.dev/en/latest/shardy_jax_migration.html)
    for more information.
  * JAX autodiff is switching to using direct linearization by default (instead of
    implementing linearization via JVP and partial eval).
    See [migration guide](https://docs.jax.dev/en/latest/direct_linearize_migration.html)
    for more information.
  * `jax.stages.OutInfo` has been replaced with `jax.ShapeDtypeStruct`.
  * `jax.jit` now requires `fun` to be passed by position, and additional
    arguments to be passed by keyword. Doing otherwise will result in an error
    starting in v0.7.x. This raised a DeprecationWarning in v0.6.x.
  * The minimum Python version is now 3.11. 3.11 will remain the minimum
    supported version until July 2026.
  * Layout API renames:
    * `Layout`, `.layout`, `.input_layouts` and `.output_layouts` have been
      renamed to `Format`, `.format`, `.input_formats` and `.output_formats`
    * `DeviceLocalLayout`, `.device_local_layout` have been renamed to `Layout`
      and `.layout`
  * `jax.experimental.shard` module has been deleted and all the APIs have been
    moved to the `jax.sharding` endpoint. So use `jax.sharding.reshard`,
    `jax.sharding.auto_axes` and `jax.sharding.explicit_axes` instead of their
    experimental endpoints.
  * `lax.infeed` and `lax.outfeed` were removed, after being deprecated in
    JAX 0.6. The `transfer_to_infeed` and `transfer_from_outfeed` methods were
    also removed the `Device` objects.
  * The `jax.extend.core.primitives.pjit_p` primitive has been renamed to
    `jit_p`, and its `name` attribute has changed from `"pjit"` to `"jit"`.
    This affects the string representations of jaxprs. The same primitive is no
    longer exported from the `jax.experimental.pjit` module.
  * The (undocumented) function `jax.extend.backend.add_clear_backends_callback`
    has been removed. Users should use `jax.extend.backend.register_backend_cache`
    instead.

* Deprecations:
  * {obj}`jax.dlpack.SUPPORTED_DTYPES` is deprecated; please use the new
    `jax.dlpack.is_supported_dtype` function.
  * `jax.scipy.special.sph_harm` has been deprecated following a similar
    deprecation in SciPy; use `jax.scipy.special.sph_harm_y` instead.
  * From {mod}`jax.interpreters.xla`, the previously deprecated symbols
    `abstractify` and `pytype_aval_mappings` have been removed.
  * `jax.interpreters.xla.canonicalize_dtype` is deprecated. For
    canonicalizing dtypes, prefer `jax.dtypes.canonicalize_dtype`.
    For checking whether an object is a valid jax input, prefer
    `jax.core.valid_jaxtype`.
  * From {mod}`jax.core`, the previously deprecated symbols `AxisName`,
    `ConcretizationTypeError`, `axis_frame`, `call_p`, `closed_call_p`,
    `get_type`, `trace_state_clean`, `typematch`, and `typecheck` have been
    removed.
  * From {mod}`jax.lib.xla_client`, the previously deprecated symbols
    `DeviceAssignment`, `get_topology_for_devices`, and `mlir_api_version`
    have been removed.
  * `jax.extend.ffi` was removed after being deprecated in v0.5.0.
    Use {mod}`jax.ffi` instead.
  * `jax.lib.xla_bridge.get_compile_options` is deprecated, and replaced by
    `jax.extend.backend.get_compile_options`.

### v0.6.2 — JAX v0.6.2

- Date: 2025-06-17
- Version: v0.6.2
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.6.2
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.6.2

- **added** — Added jax.tree.broadcast which implements a pytree prefix broadcasting helper
- **changed** — The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12

* New features:
  * Added `jax.tree.broadcast` which implements a pytree prefix broadcasting helper.

* Changes
  * The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12.

### v0.6.1 — JAX v0.6.1

- Date: 2025-05-21
- Version: v0.6.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.6.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.6.1

- **added** — Added jax.lax.axis_size which returns the size of the mapped axis given its name
- **changed** — Additional checking for the versions of CUDA package dependencies was reenabled, having been accidentally disabled in a previous release
- **changed** — jax.sharding.PartitionSpec no longer inherits from a tuple
- **changed** — jax.ShapeDtypeStruct is immutable now; use the .update method to update ShapeDtypeStruct instead of doing in-place updates
- **deprecated** — jax.custom_derivatives.custom_jvp_call_jaxpr_p is deprecated and will be removed in JAX v0.7.0

* New features:
  * Added `jax.lax.axis_size` which returns the size of the mapped axis
    given its name.

* Changes
  * Additional checking for the versions of CUDA package dependencies was
    reenabled, having been accidentally disabled in a previous release.
  * JAX nightly packages are now published to artifact registry. To install
    these packages, see the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html#jax-nightly-installation).
  * `jax.sharding.PartitionSpec` no longer inherits from a tuple.
  * `jax.ShapeDtypeStruct` is immutable now. Please use `.update` method to
    update your `ShapeDtypeStruct` instead of doing in-place updates.

* Deprecations
  * `jax.custom_derivatives.custom_jvp_call_jaxpr_p` is deprecated, and will be
    removed in JAX v0.7.0.

### v0.6.0 — JAX v0.6.0

- 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

- **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]_

### v0.5.3 — JAX v0.5.3

- Date: 2025-03-19
- Version: v0.5.3
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.5.3
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.5.3

- **added** — Added a allow_negative_indices option to jax.lax.dynamic_slice, jax.lax.dynamic_update_slice and related functions to optionally disable code generation for clamping negative indices
- **added** — Added a replace option to jax.random.categorical to enable sampling without replacement

* New Features

  * Added a `allow_negative_indices` option to `jax.lax.dynamic_slice`,
    `jax.lax.dynamic_update_slice` and related functions. The default is
    true, matching the current behavior. If set to false, JAX does not need to
    emit code clamping negative indices, which improves code size.
  * Added a `replace` option to `jax.random.categorical` to enable sampling
    without replacement.

### v0.5.2 — JAX v0.5.2

- Date: 2025-03-05
- Version: v0.5.2
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.5.2
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.5.2

- **fixed** — Fix TPU metric logging and tpu-info which was broken in 0.5.1

Patch release of 0.5.1

* Bug fixes
  * Fixes TPU metric logging and `tpu-info`, which was broken in 0.5.1

### v0.5.1 — JAX v0.5.1

- Date: 2025-02-24
- Version: v0.5.1
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.5.1
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.5.1

- **added** — Added experimental jax.experimental.custom_dce.custom_dce decorator to support customizing the behavior of opaque functions under JAX-level dead code elimination
- **added** — Added low-level reduction APIs in jax.lax: jax.lax.reduce_sum, jax.lax.reduce_prod, jax.lax.reduce_max, jax.lax.reduce_min, jax.lax.reduce_and, jax.lax.reduce_or, and jax.lax.reduce_xor
- **added** — jax.lax.linalg.qr and jax.scipy.linalg.qr now support column-pivoting on CPU and GPU
- **changed** — JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES now work as environment variables in addition to jax.config or flags
- **changed** — JAX_CPU_COLLECTIVES_IMPLEMENTATION now defaults to 'gloo', enabling multi-process CPU communication to work out-of-the-box
- **changed** — The jax[tpu] TPU extra no longer depends on the libtpu-nightly package; JAX now uses libtpu instead
- **deprecated** — The internal function linear_util.wrap_init and the constructor core.Jaxpr now require a non-empty core.DebugInfo kwarg, with a DeprecationWarning printed if jax.extend.linear_util.wrap_init is used without debugging info
- **fixed** — TPU runtime startup and shutdown time significantly improved on TPU v5e and newer, from around 17s to around 8s
- **fixed** — Persistent compilation cache no longer writes access time file if JAX_COMPILATION_CACHE_MAX_SIZE is unset or set to -1, improving performance with large-scale network storage

* New Features
  * Added an experimental `jax.experimental.custom_dce.custom_dce`
    decorator to support customizing the behavior of opaque functions under
    JAX-level dead code elimination (DCE). See `#25956` for more
    details.
  * Added low-level reduction APIs in {mod}`jax.lax`: `jax.lax.reduce_sum`,
    `jax.lax.reduce_prod`, `jax.lax.reduce_max`, `jax.lax.reduce_min`,
    `jax.lax.reduce_and`, `jax.lax.reduce_or`, and `jax.lax.reduce_xor`.
  * `jax.lax.linalg.qr`, and `jax.scipy.linalg.qr`, now support
    column-pivoting on CPU and GPU. See #20282 and
    #25955 for more details.

* Changes
  * `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` now work as
    env vars. Before they could only be specified via jax.config or flags.
  * `JAX_CPU_COLLECTIVES_IMPLEMENTATION` now defaults to `'gloo'`, meaning
    multi-process CPU communication works out-of-the-box.
  * The `jax[tpu]` TPU extra no longer depends on the `libtpu-nightly` package.
    This package may safely be removed if it is present on your machine; JAX now
    uses `libtpu` instead.

* Deprecations
  * The internal function `linear_util.wrap_init` and the constructor
    `core.Jaxpr` now must take a non-empty `core.DebugInfo` kwarg. For
    a limited time, a `DeprecationWarning` is printed if
    `jax.extend.linear_util.wrap_init` is used without debugging info.
    A downstream effect of this several other internal functions need debug
    info. This change does not affect public APIs.
    See https://github.com/jax-ml/jax/issues/26480 for more detail.

* Bug fixes
  * TPU runtime startup and shutdown time should be significantly improved on
    TPU v5e and newer (from around 17s to around 8s). If not already set, you may
    need to enable transparent hugepages in your VM image
    (`sudo sh -c 'echo always > /sys/kernel/mm/transparent_hugepage/enabled'`).
    We hope to improve this further in future releases.
  * Persistent compilation cache no longer writes access time file if
    `JAX_COMPILATION_CACHE_MAX_SIZE` is unset or set to -1, i.e. if the LRU
    eviction policy isn't enabled. This should improve performance when using
    the cache with large-scale network storage.

### v0.5.0 — JAX v0.5.0

- Date: 2025-01-17
- Version: v0.5.0
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.5.0
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.5.0

- **changed** — Enable jax_threefry_partitionable by default
- **removed** — Drop support for Mac x86 wheels
- **changed** — Increase minimum NumPy version to 1.25
- **changed** — Increase minimum SciPy version to 1.11
- **changed** — jax.numpy.einsum now defaults to optimize='auto' rather than optimize='optimal'
- **removed** — jax.numpy.linalg.solve no longer supports batched 1D arguments on the right hand side
- **added** — jax.numpy.fft.fftn, jax.numpy.fft.rfftn, jax.numpy.fft.ifftn, and jax.numpy.fft.irfftn now support transforms in more than 3 dimensions
- **added** — Add support for user defined state in the FFI via jax.ffi.register_ffi_type_id
- **added** — AOT lowering .as_text() method now supports the debug_info option to include debugging information
- **deprecated** — abstractify and pytype_aval_mappings from jax.interpreters.xla are deprecated in favor of symbols in jax.core
- **deprecated** — jax.scipy.special.lpmn and jax.scipy.special.lpmn_values are deprecated following their deprecation in SciPy v1.15.0
- **deprecated** — jax.extend.ffi submodule has been moved to jax.ffi with the previous import path deprecated
- **removed** — Delete jax_enable_memories flag
- **removed** — Remove Device and XlaRuntimeError from jax.lib.xla_client; use jax.Device and jax.errors.JaxRuntimeError instead
- **removed** — Remove jax.experimental.array_api module after being deprecated in JAX v0.4.32

As of this release, JAX now uses [effort-based versioning](https://jax.readthedocs.io/en/latest/jep/25516-effver.html).
Since this release makes a breaking change to PRNG key semantics that
may require users to update their code, we are bumping the "meso" version of JAX
to signify this.

* Breaking changes
  * Enable `jax_threefry_partitionable` by default (see
    [the update note](https://github.com/jax-ml/jax/discussions/18480)).

  * This release drops support for Mac x86 wheels. Mac ARM of course remains
    supported. For a recent discussion, see https://github.com/jax-ml/jax/discussions/22936.

    Two key factors motivated this decision:
    * The Mac x86 build (only) has a number of test failures and crashes. We
      would prefer to ship no release than a broken release.
    * Mac x86 hardware is end-of-life and cannot be easily obtained for
      developers at this point. So it is difficult for us to fix this kind of
      problem even if we wanted to.

    We are open to readding support for Mac x86 if the community is willing
    to help support that platform: in particular, we would need the JAX test
    suite to pass cleanly on Mac x86 before we could ship releases again.

* Changes:
  * The minimum NumPy version is now 1.25. NumPy 1.25 will remain the minimum
    supported version until June 2025.
  * The minimum SciPy version is now 1.11. SciPy 1.11 will remain the minimum
    supported version until June 2025.
  * `jax.numpy.einsum` now defaults to `optimize='auto'` rather than
    `optimize='optimal'`. This avoids exponentially-scaling trace-time in
    the case of many arguments (`#25214`).
  * `jax.numpy.linalg.solve` no longer supports batched 1D arguments
    on the right hand side. To recover the previous behavior in these cases,
    use `solve(a, b[..., None]).squeeze(-1)`.

* New Features
  * `jax.numpy.fft.fftn`, `jax.numpy.fft.rfftn`,
    `jax.numpy.fft.ifftn`, and `jax.numpy.fft.irfftn` now support
    transforms in more than 3 dimensions, which was previously the limit. See
    `#25606` for more details.
  * Support added for user defined state in the FFI via the new
    `jax.ffi.register_ffi_type_id` function.
  * The AOT lowering `.as_text()` method now supports the `debug_info` option
    to include debugging information, e.g., source location, in the output.

* Deprecations
  * From `jax.interpreters.xla`, `abstractify` and `pytype_aval_mappings`
    are now deprecated, having been replaced by symbols of the same name
    in `jax.core`.
  * `jax.scipy.special.lpmn` and `jax.scipy.special.lpmn_values`
    are deprecated, following their deprecation in SciPy v1.15.0. There are
    no plans to replace these deprecated functions with new APIs.
  * The `jax.extend.ffi` submodule was moved to `jax.ffi`, and the
    previous import path is deprecated.

* Deletions
  * `jax_enable_memories` flag has been deleted and the behavior of that flag
    is on by default.
  * From `jax.lib.xla_client`, the previously-deprecated `Device` and
    `XlaRuntimeError` symbols have been removed; instead use `jax.Device`
    and `jax.errors.JaxRuntimeError` respectively.
  * The `jax.experimental.array_api` module has been removed after being
    deprecated in JAX v0.4.32. Since that release, `jax.numpy` supports
    the array API directly.

### v0.4.38 — JAX v0.4.38

- Date: 2024-12-17
- Version: v0.4.38
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.38
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.4.38

- **added** — Add jax.tree.flatten_with_path and jax.tree.map_with_path as shortcuts of the corresponding tree_util functions
- **added** — Add jax.lax.split as a primitive version of jax.numpy.split that yields a more compact transpose during automatic differentiation
- **changed** — jax.export.export can be used for device-polymorphic export with shardings constructed with jax.sharding.AbstractMesh
- **deprecated** — Deprecate a number of APIs in the internal jax.core namespace
- **removed** — Remove jax.core.check_eqn, jax.core.check_type, jax.core.check_valid_jaxtype, and jax.core.non_negative_dim
- **removed** — Remove jax.lib.xla_bridge.xla_client and jax.lib.xla_bridge.default_backend
- **removed** — Remove jax.lib.xla_client._xla and jax.lib.xla_client.bfloat16
- **removed** — Remove jax.numpy.round_

* Changes:
  * `jax.tree.flatten_with_path` and `jax.tree.map_with_path` are added
    as shortcuts of the corresponding `tree_util` functions.

* Deprecations
  * a number of APIs in the internal `jax.core` namespace have been deprecated.
    Most were no-ops, were little-used, or can be replaced by APIs of the same
    name in `jax.extend.core`; see the documentation for {mod}`jax.extend`
    for information on the compatibility guarantees of these semi-public extensions.
  * Several previously-deprecated APIs have been removed, including:
    * from `jax.core`: `check_eqn`, `check_type`,  `check_valid_jaxtype`, and
      `non_negative_dim`.
    * from `jax.lib.xla_bridge`: `xla_client` and `default_backend`.
    * from `jax.lib.xla_client`: `_xla` and `bfloat16`.
    * from `jax.numpy`: `round_`.

* New Features
  * `jax.export.export` can be used for device-polymorphic export with
    shardings constructed with {func}`jax.sharding.AbstractMesh`.
    See the [jax.export documentation](https://jax.readthedocs.io/en/latest/export/export.html#device-polymorphic-export).
  * Added `jax.lax.split`. This is a primitive version of
    `jax.numpy.split`, added because it yields a more compact
    transpose during automatic differentiation.

### v0.4.37 — JAX v0.4.37

- Date: 2024-12-10
- Version: v0.4.37
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.37
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.4.37

- **fixed** — Fixed a bug where jit would error if an argument was named f
- **fixed** — Fixed a bug that throws index out of range error in jax.lax.while_loop if the user registers pytree node class with different aux data for the flatten and flatten_with_path
- **fixed** — Pinned libtpu release 0.0.6 that fixes a compiler bug on TPU v6e

This is a patch release of jax 0.4.36. Only "jax" was released at this version.

* Bug fixes
  * Fixed a bug where `jit` would error if an argument was named `f` (#25329).
  * Fix a bug that will throw `index out of range` error in
    `jax.lax.while_loop` if the user registers pytree node class with
    different aux data for the flatten and flatten_with_path.
  * Pinned a new libtpu release (0.0.6) that fixes a compiler bug on TPU v6e.

### v0.4.36 — JAX v0.4.36

- Date: 2024-12-05
- Version: v0.4.36
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.36
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.4.36

- **changed** — JAX tracing machinery now uses stackless dispatch as a function of context rather than both context and data, removing levels, sublevels, post_process_call, new_base_main, and custom_bind
- **removed** — Removed support for jax.experimental.jax2tf.convert with native_serialization=False or enable_xla=False
- **removed** — Removed xb, xc, and xe symbols from jax.interpreters.xla
- **removed** — Removed deprecated module jax.experimental.export
- **removed** — Removed initial argument to jax.nn.softmax and jax.nn.log_softmax
- **changed** — Calling np.asarray on typed PRNG keys now raises an error instead of returning a scalar object array
- **removed** — Removed jax.export.DisabledSafetyCheck.shape_assertions, jax.export.Exported.lowering_platforms, jax.export.Exported.mlir_module_serialization_version, jax.export.Exported.uses_shape_polymorphism, and lowering_platforms kwarg from jax.export.export
- **removed** — Removed symbolic_scope and symbolic_constraints kwargs from jax.export.symbolic_args_specs
- **changed** — Hashing of tracers now raises a TypeError
- **changed** — JAX build CLI now uses a subcommand structure with build and requirements_update subcommands
- **changed** — jax.scipy.linalg.toeplitz now does implicit batching on multi-dimensional inputs
- **changed** — jax.scipy.special.gamma and jax.scipy.special.gammasgn now return NaN for negative integer inputs
- **removed** — Removed jax.clear_backends
- **removed** — Removed __gpu$xla.gpu.triton custom call from export stability guarantees
- **added** — Added compiler_options argument to jax.jit for passing compilation options to XLA
- **changed** — jax.tree_util.register_dataclass now allows metadata fields to be declared inline via dataclasses.field
- **added** — Added jax.numpy.put_along_axis
- **added** — Added GPU support for jax.lax.linalg.eig and related jax.numpy functions jax.numpy.linalg.eig and jax.numpy.linalg.eigvals
- **added** — Added jax_exec_time_optimization_effort and jax_memory_fitting_effort configuration flags to control compiler effort for execution time and memory usage
- **fixed** — Fixed indexing overflow bug in GPU implementations of LU and QR decomposition for batch sizes close to int32 max
- **deprecated** — Deprecated jax.lib.xla_extension.ArrayImpl and jax.lib.xla_client.ArrayImpl in favor of jax.Array
- **deprecated** — Deprecated jax.lib.xla_extension.XlaRuntimeError in favor of jax.errors.JaxRuntimeError

* Breaking Changes
  * This release lands "stackless", an internal change to JAX's tracing
    machinery. We made trace dispatch purely a function of context rather than a
    function of both context and data. This let us delete a lot of machinery for
    managing data-dependent tracing: levels, sublevels, `post_process_call`,
    `new_base_main`, `custom_bind`, and so on. The change should only affect
    users that use JAX internals.

    If you do use JAX internals then you may need to
    update your code (see
    https://github.com/jax-ml/jax/commit/c36e1f7c1ad4782060cbc8e8c596d85dfb83986f
    for clues about how to do this). There might also be version skew
    issues with JAX libraries that do this. If you find this change breaks your
    non-JAX-internals-using code then try the
    `config.jax_data_dependent_tracing_fallback` flag as a workaround, and if
    you need help updating your code then please file a bug.
  * `jax.experimental.jax2tf.convert` with `native_serialization=False`
    or with `enable_xla=False` have been deprecated since July 2024, with
    JAX version 0.4.31. Now we removed support for these use cases. `jax2tf`
    with native serialization will still be supported.
  * In `jax.interpreters.xla`, the `xb`, `xc`, and `xe` symbols have been removed
    after being deprecated in JAX v0.4.31. Instead use `xb = jax.lib.xla_bridge`,
    `xc = jax.lib.xla_client`, and `xe = jax.lib.xla_extension`.
  * The deprecated module `jax.experimental.export` has been removed. It was replaced
    by `jax.export` in JAX v0.4.30. See the [migration guide](https://jax.readthedocs.io/en/latest/export/export.html#migration-guide-from-jax-experimental-export)
    for information on migrating to the new API.
  * The `initial` argument to `jax.nn.softmax` and `jax.nn.log_softmax`
    has been removed, after being deprecated in v0.4.27.
  * Calling `np.asarray` on typed PRNG keys (i.e. keys produced by `jax.random.key`)
    now raises an error. Previously, this returned a scalar object array.
  * The following deprecated methods and functions in `jax.export` have
    been removed:
      * `jax.export.DisabledSafetyCheck.shape_assertions`: it had no effect
        already.
      * `jax.export.Exported.lowering_platforms`: use `platforms`.
      * `jax.export.Exported.mlir_module_serialization_version`:
        use `calling_convention_version`.
      * `jax.export.Exported.uses_shape_polymorphism`:
         use `uses_global_constants`.
      * the `lowering_platforms` kwarg for `jax.export.export`: use
        `platforms` instead.
  * The kwargs `symbolic_scope` and `symbolic_constraints` from
    `jax.export.symbolic_args_specs` have been removed. They were
    deprecated in June 2024. Use `scope` and `constraints` instead.
  * Hashing of tracers, which has been deprecated since version 0.4.30, now
    results in a `TypeError`.
  * Refactor: JAX build CLI (build/build.py) now uses a subcommand structure and
    replaces previous build.py usage. Run `python build/build.py --help` for
    more details. Brief overview of the new subcommand options:
    * `build`: Builds JAX wheel packages. For e.g., `python build/build.py build --wheels=jaxlib,jax-cuda-pjrt`
    * `requirements_update`: Updates requirements_lock.txt files.
  * `jax.scipy.linalg.toeplitz` now does implicit batching on multi-dimensional
    inputs. To recover the previous behavior, you can call `jax.numpy.ravel`
    on the function inputs.
  * `jax.scipy.special.gamma` and `jax.scipy.special.gammasgn` now
    return NaN for negative integer inputs, to match the behavior of SciPy from
    https://github.com/scipy/scipy/pull/21827.
  * `jax.clear_backends` was removed after being deprecated in v0.4.26.
  * We removed the custom call "__gpu$xla.gpu.triton" from the list of custom
    call that we guarantee export stability. This is because this custom call
    relies on Triton IR, which is not guaranteed t

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