jaxlib

• New features

  • Added an error check for trying to deserialize JAX exports that are older
    than the backwards compatibility window. Without this check the
    deserialization of expired artifacts may succeed and then result in
    obscure downstream errors.
    Added a configuration flag --jax_export_deserialize_expired_versions to
    temporarily bypass the error check.
    See https://docs.jax.dev/en/latest/export/export.html#compatibility-guarantees.
  • Added jax.numpy.top_k, which implements numpy.top_k, added in
    in NumPy v2.6.0 (#39729).

• Breaking changes

  • The exec_time_optimization_effort and memory_fitting_effort flags have been
    removed in favor of the EffortLevel enum.
  • JAX does not support anymore deserialization of Exported modules from
    before January 15th, 2026 because they are beyond the backwards compatibility
    window. On that date we added support to serialize shardings as NamedSharding,
    and now that is the only sharding serialization that is supported.
  • jnp.take_along_axis now always defaults wrap_negative_indices to True.
    It used to default to False for mode=promise_in_bounds and True otherwise.
    (This also means None is no longer a valid value for wrap_negative_indices.)

• Deprecations

  • The fields in_shardings_hlo and out_shardings_hlo of
    jax.export.Exported have been deprecated for a while. Now accessing them
    raises a warning. Use in_shardings_jax and out_shardings_jax instead.

• Changes

  • The cuDNN fused attention backward pass (used by
    jax.nn.dot_product_attention with implementation='cudnn') no
    longer computes a bias gradient when the only attention bias comes from
    a boolean mask, whose gradient no caller can request. Bias gradients
    for an explicit bias or a non-boolean mask are unchanged
    (#34685).
  • jax.numpy.meshgrid, jax.numpy.ogrid, and
    jax.numpy.broadcast_arrays now return tuples rather than lists
    in order to align with NumPy>2.0 and the Array API specification.
    (#39783, #39789, #39802)
  • When jax.grad or jax.value_and_grad rejects a function with
    a non-scalar output, the error message now suggests reducing the output
    to a scalar (e.g. with output.sum()), using jax.jacobian, or
    reshaping size-1 outputs (#2303).
  • When indexing with non-static or traced slice indices, the error message
    now suggests using jax.lax.dynamic_slice,
    jax.lax.dynamic_update_slice, or jax.ds, and shows tracer
    provenance (#7222).
  • PyTree metadata equality comparison failures now report the specific
    registered pytree type that caused the error (#13027).

• Bug fixes

  • jax.numpy.linalg.det and jax.numpy.linalg.slogdet now use a
    closed-form LU decomposition with row pivoting for 2x2 and 3x3 matrices
    instead of closed-form polynomial expansions to avoid numerical instability
    and catastrophic cancellation (#39905).
  • The batching rules of the cuDNN fused attention primitives (used by
    jax.nn.dot_product_attention with implementation='cudnn') now
    support operands that do not carry the vmap axis, including a shared
    bias or mask. Previously jax.jacobian, jax.vmap with partial
    in_axes, and jax.vmap of a VJP or of jax.grad failed with a
    reshape TypeError (#38495).
  • jax.vmap of fp8 cuDNN fused attention now works: its batching rules
    additionally mislabeled or dropped the amax outputs and restored output
    shapes incorrectly, so previously no vmap of the fp8 path succeeded at
    all. The amax outputs are whole-batch statistics and do not carry the
    vmap axis; vmap over the scale/descale operands raises a clear
    NotImplementedError.
  • Setting jax_compiler_enable_remat_pass to False now adds
    rematerialization to the set of disabled XLA passes instead of
    overwriting it, so HLO passes disabled via
    XLA_FLAGS=--xla_disable_hlo_passes=... stay disabled
    (#37391).
  • jax.numpy.split, jax.numpy.array_split, and the
    hsplit/vsplit/dsplit variants once again accept negative entries in
    indices_or_sections, resolving them against the axis size as NumPy does
    (#6599). Out-of-bound indices are now clipped to the axis
    bounds and produce empty sections, also matching NumPy, instead of raising
    ValueError: Sizes passed to split must be nonnegative.
  • Fixed abstract evaluation in jax.lax.scan to only check .mat
    equivalency when the abstract value is a ShapedArray
    (#39700).
  • Fixed propagation of singleton sharded dimensions in jax.lax.reshape
    when reshaping arrays with sharding constraints (#39309).
  • Fixed jax.tree_util.flatten_one_level_with_keys for namedtuple
    instances (#39297).
  • Fixed _get_prime_factors in jax.experimental.mesh_utils
    (#38286).

• 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 jax.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.
  • 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.

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

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

• 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.experimental.custom_dce.custom_dce
  • jax.scipy.linalg.cho_solve, jax.scipy.linalg.lu_solve, and
    jax.scipy.linalg.solve_triangular now show a deprecation warning for
    batched 1D solves with b.ndim > 1. In the future these will be treated as
    batched 2D solves.
  • Added a new version 10 for the jax.export serialization format. This is
    an optimization for when there are multiple occurrences of the same
    abstract value, abstract mesh, or sharding.

• Bug fixes:

  • Fixed a bug that led to differing output between CPU and GPU for
    non-symmetric multidimensional IRFFTs (#29325).
  • Fixed an error when tiny matrices were passed to
    jax.lax.linalg.tridiagonal_solve on GPU (#32487).
  • Fixed a bug in jax.scipy.fft.dctn and idctn where axes=None
    incorrectly defaulted to all axes when s was specified, instead of the
    last len(s) axes to match SciPy behavior (#29426).
  • Fixed a bug where calling jax.distributed.initialize() on a GCE TPU
    Managed Instance Group raised an IndexError (#36593). When
    jax.distributed.initialize() is called on a GCE VM, it uses the GCE
    metadata server to learn the addresses of all participating tasks. The format of this metadata
    on Managed Instance Groups was not a format JAX expected, leading to the
    exception. We now parse this format correctly.

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(...)).

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

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

• openxla/xla#36579
• openxla/xla#36345
• openxla/xla#36755
• openxla/xla#36696

• 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.pmapped 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.

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

• openxla/xla@4bc723d
• openxla/xla@21552fd

• Deprecations

  • jax.lax.pvary has been deprecated. Please use jax.lax.pcast(..., to='varying') as the replacement.
  • Complex arguments passed to {func}jax.numpy.arange now result in a deprecation warning, because the output is poorly-defined.
  • From {mod}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 {mod}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 Tracers represent Arrays, 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.

• New features:

  • {func}jax.jit now supports the decorator factory pattern; i.e instead of writing

    @functools.partial(jax.jit, static_argnames=['n'])
    def f(x, n):
      ...
    

    you may write

    @jax.jit(static_argnames=['n'])
    def f(x, n):
      ...
    

• Changes:

  • {func}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 {mod}jax.lax.linalg.

  • {func}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 (({jax-issue}#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.
  • {mod}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.

• 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.
  • {func}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 {func}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 {obj}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 {mod}jax.experimental.compilation_cache, the deprecated symbols is_initialized and initialize_cache were removed.
  • The deprecated function {func}jax.interpreters.xla.canonicalize_dtype was removed.
  • {mod}jaxlib.hlo_helpers has been removed. Use {mod}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 {func}jax.numpy.percentile and {func}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 {func}jax.lax.fori_loop. If you need help updating your code, please file a bug.
  • {func}jax.numpy.trimzeros now errors for non-1D input.
  • The where argument to {func}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 {mod} jax.dlpack, {mod} jax.errors, {mod} jax.lib.xla_bridge, {mod} jax.lib.xla_client, and {mod} jax.lib.xla_extension were removed.
  • jax.interpreters.mlir.dense_bool_array was removed. Use MLIR APIs to construct attributes instead.

• Changes

  • {func}jax.numpy.linalg.eig now returns a namedtuple (with attributes eigenvalues and eigenvectors) instead of a plain tuple.
  • {func}jax.grad and {func}jax.vjp will now round always primals to float32 if float64 mode is not enabled.
  • {func}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 {func}jax.lax.linalg.eig ({jax-issue}#27265). The use_magma argument is now deprecated in favor of implementation.
  • {func}jax.numpy.trim_zeros now follows NumPy 2.2 in supporting multi-dimensional inputs.

• Deprecations

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

• Breaking changes:

  • {func}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 TypedNdArray, 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.

• 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.
  • {func}jax.lax.dot now implements the general dot product via the optional dimension_numbers argument.

• Deprecations:

  • {func}jax.lax.zeros_like_array is deprecated. Please use {func}jax.numpy.zeros_like instead.
  • Attempting to import {mod}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 {func}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 {mod}jax.interpreters.ad, {mod}jax.interpreters.batching, and {mod}jax.interpreters.partial_eval; they are used rarely if ever outside JAX itself, and most are deprecated without any public replacement.

• New features:

  • Added jax.P which is an alias for jax.sharding.PartitionSpec.
  • Added {func}jax.tree.reduce_associative.
  • The {attr}jax.numpy.ndarray.at indexing methods now support a wrap_negative_indices argument, which defaults to True to match the current behavior ({jax-issue}#29434).

• Breaking changes:

  • JAX is migrating from GSPMD to Shardy by default. See the migration guide 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 for more information.
  • jax.stages.OutInfo has been replaced with jax.ShapeDtypeStruct.
  • {func}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.
  • out_sharding arg added to x.at[y].set and x.at[y].add. Previous behavior propagating operand sharding removed. Please use x.at[y].set/add(z, out_sharding=jax.typeof(x).sharding) to retain previous behavior if scatter op requires collectives.

• Deprecations:

  • {obj}jax.dlpack.SUPPORTED_DTYPES is deprecated; please use the new {func}jax.dlpack.is_supported_dtype function.
  • {func}jax.scipy.special.sph_harm has been deprecated following a similar deprecation in SciPy; use {func}jax.scipy.special.sph_harm_y instead.
  • From {mod}jax.interpreters.xla, the previously deprecated symbols abstractify and pytype_aval_mappings have been removed.
  • {func}jax.interpreters.xla.canonicalize_dtype is deprecated. For canonicalizing dtypes, prefer {func}jax.dtypes.canonicalize_dtype. For checking whether an object is a valid jax input, prefer {func}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.
  • {func}jax.lib.xla_bridge.get_compile_options is deprecated, and replaced by {func}jax.extend.backend.get_compile_options.

• New features:

  • Added {func}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.

• New features:

  • Added {func}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 re-enabled, 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.
  • 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.

• Breaking changes

  • {func}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 {jax-issue}#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.
  • {func}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

  • {func}jax.tree_util.build_tree is deprecated. Use {func}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 {mod}jax.ffi should be used instead.
  • The deprecated use of {func}jax.ffi.ffi_call with inline arguments is no longer supported. {func}~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 {mod}jax.tree or {mod}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 {mod}jax.extend.core.
    • The vectorized argument to {func}~jax.pure_callback and {func}~jax.ffi.ffi_call. Use the vmap_method parameter instead.

• New Features

  • Added a allow_negative_indices option to {func}jax.lax.dynamic_slice, {func}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 {func}jax.random.categorical to enable sampling without replacement.

• Breaking changes

  • The jit tracing cache now keys on input NamedShardings. Previously, the tracing cache did not include sharding information at all (although subsequent jit caches did like lowering and compilation caches), so two equivalent shardings of different types would not retrace, but now they do. For example:

    @jax.jit
    def f(x):
      return x
    
    # inp1.sharding is of type SingleDeviceSharding
    inp1 = jnp.arange(8)
    f(inp1)
    
    mesh = jax.make_mesh((1,), ('x',))
    # inp2.sharding is of type NamedSharding
    inp2 = jax.device_put(jnp.arange(8), NamedSharding(mesh, P('x')))
    f(inp2)  # tracing cache miss
    

    In the above example, calling f(inp1) and then f(inp2) will lead to a tracing cache miss because the shardings have changed on the abstract values while tracing.

• New Features

  • Added an experimental {func}jax.experimental.custom_dce.custom_dce decorator to support customizing the behavior of opaque functions under JAX-level dead code elimination (DCE). See {jax-issue}#25956 for more details.
  • Added low-level reduction APIs in {mod}jax.lax: {func}jax.lax.reduce_sum, {func}jax.lax.reduce_prod, {func}jax.lax.reduce_max, {func}jax.lax.reduce_min, {func}jax.lax.reduce_and, {func}jax.lax.reduce_or, and {func}jax.lax.reduce_xor.
  • {func}jax.lax.linalg.qr, and {func}jax.scipy.linalg.qr, now support column-pivoting on CPU and GPU. See {jax-issue}#20282 and
  • Added {func}jax.random.multinomial. {jax-issue}#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.
  • In {func}jax.numpy.ndim, {func}jax.numpy.shape, and {func}jax.numpy.size, non-arraylike inputs (such as lists, tuples, etc.) are now deprecated.

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

As of this release, JAX now uses effort-based versioning. 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).

  • 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 re-adding 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.
  • {func}jax.numpy.einsum now defaults to optimize='auto' rather than optimize='optimal'. This avoids exponentially-scaling trace-time in the case of many arguments ({jax-issue}#25214).
  • {func}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

  • {func}jax.numpy.fft.fftn, {func}jax.numpy.fft.rfftn, {func}jax.numpy.fft.ifftn, and {func}jax.numpy.fft.irfftn now support transforms in more than 3 dimensions, which was previously the limit. See {jax-issue}#25606 for more details.
  • Support added for user defined state in the FFI via the new {func}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 {mod}jax.interpreters.xla, abstractify and pytype_aval_mappings are now deprecated, having been replaced by symbols of the same name in {mod}jax.core.
  • {func}jax.scipy.special.lpmn and {func}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 {mod}jax.extend.ffi submodule was moved to {mod}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, {mod}jax.numpy supports the array API directly.

• Breaking Changes

  • XlaExecutable.cost_analysis now returns a dict[str, float] (instead of a single-element list[dict[str, float]]).

• 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 {mod}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 {mod}jax.core: check_eqn, check_type, check_valid_jaxtype, and non_negative_dim.
    • from {mod}jax.lib.xla_bridge: xla_client and default_backend.
    • from {mod}jax.lib.xla_client: _xla and bfloat16.
    • from {mod}jax.numpy: round_.

• New Features

  • {func}jax.export.export can be used for device-polymorphic export with shardings constructed with {func}jax.sharding.AbstractMesh. See the jax.export documentation.
  • Added {func}jax.lax.split. This is a primitive version of {func}jax.numpy.split, added because it yields a more compact transpose during automatic differentiation.

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

  • {func}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 {mod}jax.export in JAX v0.4.30. See the migration guide for information on migrating to the new API.

  • The initial argument to {func}jax.nn.softmax and {func}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 {func}jax.random.key) now raises an error. Previously, this returned a scalar object array.

  • The following deprecated methods and functions in {mod}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 {func}jax.export.export: use platforms instead.

  • The kwargs symbolic_scope and symbolic_constraints from {func}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.

  • {func}jax.scipy.linalg.toeplitz now does implicit batching on multi-dimensional inputs. To recover the previous behavior, you can call {func}jax.numpy.ravel on the function inputs.

  • {func}jax.scipy.special.gamma and {func}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 to be stable. If you need to export code that uses this custom call, you can use the disabled_checks parameter. See more details in the documentation.

• New Features

  • {func}jax.jit got a new compiler_options: dict[str, Any] argument, for passing compilation options to XLA. For the moment it's undocumented and may be in flux.
  • {func}jax.tree_util.register_dataclass now allows metadata fields to be declared inline via {func}dataclasses.field. See the function documentation for examples.
  • Added {func}jax.numpy.put_along_axis.
  • {func}jax.lax.linalg.eig and the related jax.numpy functions ({func}jax.numpy.linalg.eig and {func}jax.numpy.linalg.eigvals) are now supported on GPU. See {jax-issue}#24663 for more details.
  • Added two new configuration flags, jax_exec_time_optimization_effort and jax_memory_fitting_effort, to control the amount of effort the compiler spends minimizing execution time and memory usage, respectively. Valid values are between -1.0 and 1.0, default is 0.0.
  • Added {func}jax.scipy.special.loggamma

• Bug fixes

  • Fixed a bug where the GPU implementations of LU and QR decomposition would result in an indexing overflow for batch sizes close to int32 max. See {jax-issue}#24843 for more details.

• Deprecations

  • jax.lib.xla_extension.ArrayImpl and jax.lib.xla_client.ArrayImpl are deprecated; use jax.Array instead.
  • jax.lib.xla_extension.XlaRuntimeError is deprecated; use jax.errors.JaxRuntimeError instead.

• Breaking Changes

  • {func}jax.numpy.isscalar now returns True for any array-like object with zero dimensions. Previously it only returned True for zero-dimensional array-like objects with a weak dtype.
  • jax.experimental.host_callback has been deprecated since March 2024, with JAX version 0.4.26. Now we removed it. See {jax-issue}#20385 for a discussion of alternatives.

• Changes:

  • jax.lax.FftType was introduced as a public name for the enum of FFT operations. The semi-public API jax.lib.xla_client.FftType has been deprecated.
  • TPU: JAX now installs TPU support from the libtpu package rather than libtpu-nightly. For the next few releases JAX will pin an empty version of libtpu-nightly as well as libtpu to ease the transition; that dependency will be removed in Q1 2025.

• Deprecations:

  • The semi-public API jax.lib.xla_client.PaddingType has been deprecated. No JAX APIs consume this type, so there is no replacement.
  • The default behavior of {func}jax.pure_callback and {func}jax.extend.ffi.ffi_call under vmap has been deprecated and so has the vectorized parameter to those functions. The vmap_method parameter should be used instead for better defined behavior. See the discussion in {jax-issue}#23881 for more details.
  • The semi-public API jax.lib.xla_client.register_custom_call_target has been deprecated. Use the JAX FFI instead.
  • The semi-public APIs jax.lib.xla_client.dtype_to_etype, jax.lib.xla_client.ops, jax.lib.xla_client.shape_from_pyval, jax.lib.xla_client.PrimitiveType, jax.lib.xla_client.Shape, jax.lib.xla_client.XlaBuilder, and jax.lib.xla_client.XlaComputation have been deprecated. Use StableHLO instead.

• New Functionality

  • This release includes wheels for Python 3.13. Free-threading mode is not yet supported.
  • jax.errors.JaxRuntimeError has been added as a public alias for the formerly private XlaRuntimeError type.

• Breaking changes

  • jax_pmap_no_rank_reduction flag is set to True by default.
    • array[0] on a pmap result now introduces a reshape (use array[0:1] instead).
    • The per-shard shape (accessible via jax_array.addressable_shards or jax_array.addressable_data(0)) now has a leading (1, ...). Update code that directly accesses shards accordingly. The rank of the per-shard-shape now matches that of the global shape which is the same behavior as jit. This avoids costly reshapes when passing results from pmap into jit.
  • jax.experimental.host_callback has been deprecated since March 2024, with JAX version 0.4.26. Now we set the default value of the --jax_host_callback_legacy configuration value to True, which means that if your code uses jax.experimental.host_callback APIs, those API calls will be implemented in terms of the new jax.experimental.io_callback API. If this breaks your code, for a very limited time, you can set the --jax_host_callback_legacy to True. Soon we will remove that configuration option, so you should instead transition to using the new JAX callback APIs. See {jax-issue}#20385 for a discussion.

• Deprecations

  • In {func}jax.numpy.trim_zeros, non-arraylike arguments or arraylike arguments with ndim != 1 are now deprecated, and in the future will result in an error.
  • Internal pretty-printing tools jax.core.pp_* have been removed, after being deprecated in JAX v0.4.30.
  • jax.lib.xla_client.Device is deprecated; use jax.Device instead.
  • jax.lib.xla_client.XlaRuntimeError has been deprecated. Use jax.errors.JaxRuntimeError instead.

• Deletion:

  • jax.xla_computation is deleted. It's been 3 months since it's deprecation in 0.4.30 JAX release. Please use the AOT APIs to get the same functionality as jax.xla_computation.
    • jax.xla_computation(fn)(*args, **kwargs) can be replaced with jax.jit(fn).lower(*args, **kwargs).compiler_ir('hlo').
    • You can also use .out_info property of jax.stages.Lowered to get the output information (like tree structure, shape and dtype).
    • For cross-backend lowering, you can replace jax.xla_computation(fn, backend='tpu')(*args, **kwargs) with jax.jit(fn).trace(*args, **kwargs).lower(lowering_platforms=('tpu',)).compiler_ir('hlo').
  • {class}jax.ShapeDtypeStruct no longer accepts the named_shape argument. The argument was only used by xmap which was removed in 0.4.31.
  • jax.tree.map(f, None, non-None), which previously emitted a DeprecationWarning, now raises an error in a future version of jax. None is only a tree-prefix of itself. To preserve the current behavior, you can ask jax.tree.map to treat None as a leaf value by writing: jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None).
  • jax.sharding.XLACompatibleSharding has been removed. Please use jax.sharding.Sharding.

• Bug fixes

  • Fixed a bug where {func}jax.numpy.cumsum would produce incorrect outputs if a non-boolean input was provided and dtype=bool was specified.
  • Edit implementation of {func}jax.numpy.ldexp to get correct gradient.

This is a patch release on top of jax 0.4.32, that fixes two bugs found in that release.

A TPU-only data corruption bug was found in the version of libtpu pinned by JAX 0.4.32, which manifested only if multiple TPU slices were present in the same job, for example, if training on multiple v5e slices. This release fixes that issue by pinning a fixed version of libtpu.

This release fixes an inaccurate result for F64 tanh on CPU (#23590).

WARNING: This release has been yanked from PyPI because of a data corruption bug on TPU if there are multiple TPU slices in the job

• Deletion

  • xmap has been deleted. Please use {func}shard_map as the replacement.

• Changes

  • The minimum CuDNN version is v9.1. This was true in previous releases also, but we now declare this version constraint formally.
  • The minimum Python version is now 3.10. 3.10 will remain the minimum supported version until July 2025.
  • The minimum NumPy version is now 1.24. NumPy 1.24 will remain the minimum supported version until December 2024.
  • The minimum SciPy version is now 1.10. SciPy 1.10 will remain the minimum supported version until January 2025.
  • {func}jax.numpy.ceil, {func}jax.numpy.floor and {func}jax.numpy.trunc now return the output of the same dtype as the input, i.e. no longer upcast integer or boolean inputs to floating point.
  • libdevice.10.bc is no longer bundled with CUDA wheels. It must be installed either as a part of local CUDA installation, or via NVIDIA's CUDA pip wheels.
  • {class}jax.experimental.pallas.BlockSpec now expects block_shape to be passed before index_map. The old argument order is deprecated and will be removed in a future release.
  • Updated the repr of gpu devices to be more consistent with TPUs/CPUs. For example, cuda(id=0) will now be CudaDevice(id=0).
  • Added the device property and to_device method to {class}jax.Array, as part of JAX's Array API support.

• Deprecations

  • Removed a number of previously-deprecated internal APIs related to polymorphic shapes. From {mod}jax.core: removed canonicalize_shape, dimension_as_value, definitely_equal, and symbolic_equal_dim.
  • HLO lowering rules should no longer wrap singleton ir.Values in tuples. Instead, return singleton ir.Values unwrapped. Support for wrapped values will be removed in a future version of JAX.
  • {func}jax.experimental.jax2tf.convert with native_serialization=False or enable_xla=False is now deprecated and this support will be removed in a future version. Native serialization has been the default since JAX 0.4.16 (September 2023).
  • The previously-deprecated function jax.random.shuffle has been removed; instead use jax.random.permutation with independent=True.

• Changes

  • JAX supports ml_dtypes >= 0.2. In 0.4.29 release, the ml_dtypes version was bumped to 0.4.0 but this has been rolled back in this release to give users of both TensorFlow and JAX more time to migrate to a newer TensorFlow release.
  • jax.experimental.mesh_utils can now create an efficient mesh for TPU v5e.
  • jax now depends on jaxlib directly. This change was enabled by the CUDA plugin switch: there are no longer multiple jaxlib variants. You can install a CPU-only jax with pip install jax, no extras required.
  • Added an API for exporting and serializing JAX functions. This used to exist in jax.experimental.export (which is being deprecated), and will now live in jax.export. See the documentation.

• Deprecations

  • Internal pretty-printing tools jax.core.pp_* are deprecated, and will be removed in a future release.
  • Hashing of tracers is deprecated, and will lead to a TypeError in a future JAX release. This previously was the case, but there was an inadvertent regression in the last several JAX releases.
  • jax.experimental.export is deprecated. Use {mod}jax.export instead. See the migration guide.
  • Passing an array in place of a dtype is now deprecated in most cases; e.g. for arrays x and y, x.astype(y) will raise a warning. To silence it use x.astype(y.dtype).
  • jax.xla_computation is deprecated and will be removed in a future release. Please use the AOT APIs to get the same functionality as jax.xla_computation.
    • jax.xla_computation(fn)(*args, **kwargs) can be replaced with jax.jit(fn).lower(*args, **kwargs).compiler_ir('hlo').
    • You can also use .out_info property of jax.stages.Lowered to get the output information (like tree structure, shape and dtype).
    • For cross-backend lowering, you can replace jax.xla_computation(fn, backend='tpu')(*args, **kwargs) with jax.jit(fn).trace(*args, **kwargs).lower(lowering_platforms=('tpu',)).compiler_ir('hlo').

• Bug fixes

  • Fixed a bug where XLA sharded some concatenation operations incorrectly, which manifested as an incorrect output for cumulative reductions (#21403).
  • Fixed a bug where XLA:CPU miscompiled certain matmul fusions (https://github.com/openxla/xla/pull/13301).
  • Fixes a compiler crash on GPU (https://github.com/google/jax/issues/21396).

• Deprecations

  • jax.tree.map(f, None, non-None) now emits a DeprecationWarning, and will raise an error in a future version of jax. None is only a tree-prefix of itself. To preserve the current behavior, you can ask jax.tree.map to treat None as a leaf value by writing: jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None).

• Bug fixes

  • Fixes a memory corruption bug in the type name of Array and JIT Python objects in Python 3.10 or earlier.
  • Fixed a warning '+ptx84' is not a recognized feature for this target under CUDA 12.4.
  • Fixed a slow compilation problem on CPU.

• Changes

  • The Windows build is now built with Clang instead of MSVC.

• New Functionality

  • Added {func}jax.numpy.unstack and {func}jax.numpy.cumulative_sum, following their addition in the array API 2023 standard, soon to be adopted by NumPy.
  • Added a new config option jax_cpu_collectives_implementation to select the implementation of cross-process collective operations used by the CPU backend. Choices available are 'none'(default), 'gloo' and 'mpi' (requires jaxlib 0.4.26). If set to 'none', cross-process collective operations are disabled.

• Changes

  • {func}jax.pure_callback, {func}jax.experimental.io_callback and {func}jax.debug.callback now use {class}jax.Array instead of {class}np.ndarray. You can recover the old behavior by transforming the arguments via jax.tree.map(np.asarray, args) before passing them to the callback.
  • complex_arr.astype(bool) now follows the same semantics as NumPy, returning False where complex_arr is equal to 0 + 0j, and True otherwise.
  • core.Token now is a non-trivial class which wraps a jax.Array. It could be created and threaded in and out of computations to build up dependency. The singleton object core.token has been removed, users now should create and use fresh core.Token objects instead.
  • On GPU, the Threefry PRNG implementation no longer lowers to a kernel call by default. This choice can improve runtime memory usage at a compile-time cost. Prior behavior, which produces a kernel call, can be recovered with jax.config.update('jax_threefry_gpu_kernel_lowering', True). If the new default causes issues, please file a bug. Otherwise, we intend to remove this flag in a future release.

• Deprecations & Removals

  • Pallas now exclusively uses XLA for compiling kernels on GPU. The old lowering pass via Triton Python APIs has been removed and the JAX_TRITON_COMPILE_VIA_XLA environment variable no longer has any effect.
  • {func}jax.numpy.clip has a new argument signature: a, a_min, and a_max are deprecated in favor of x (positional only), min, and max ({jax-issue}20550).
  • The device() method of JAX arrays has been removed, after being deprecated since JAX v0.4.21. Use arr.devices() instead.
  • The initial argument to {func}jax.nn.softmax and {func}jax.nn.log_softmax is deprecated; empty inputs to softmax are now supported without setting this.
  • In {func}jax.jit, passing invalid static_argnums or static_argnames now leads to an error rather than a warning.
  • The minimum jaxlib version is now 0.4.23.
  • The {func}jax.numpy.hypot function now issues a deprecation warning when passing complex-valued inputs to it. This will raise an error when the deprecation is completed.
  • Scalar arguments to {func}jax.numpy.nonzero, {func}jax.numpy.where, and related functions now raise an error, following a similar change in NumPy.
  • The config option jax_cpu_enable_gloo_collectives is deprecated. Use jax.config.update('jax_cpu_collectives_implementation', 'gloo') instead.
  • The jax.Array.device_buffer and jax.Array.device_buffers methods have been removed after being deprecated in JAX v0.4.22. Instead use {attr}jax.Array.addressable_shards and {meth}jax.Array.addressable_data.
  • The condition, x, and y parameters of jax.numpy.where are now positional-only, following deprecation of the keywords in JAX v0.4.21.
  • Non-array arguments to functions in {mod}jax.lax.linalg now must be specified by keyword. Previously, this raised a DeprecationWarning.
  • Array-like arguments are now required in several {func}jax.numpy APIs, including {func}~jax.numpy.apply_along_axis, {func}~jax.numpy.apply_over_axes, {func}~jax.numpy.inner, {func}~jax.numpy.outer, {func}~jax.numpy.cross, {func}~jax.numpy.kron, and {func}~jax.numpy.lexsort.

• Bug fixes

  • {func}jax.numpy.astype will now always return a copy when copy=True. Previously, no copy would be made when the output array would have the same dtype as the input array. This may result in some increased memory usage. The default value is set to copy=False to preserve backwards compatibility.

• New Functionality

  • Added {func}jax.numpy.trapezoid, following the addition of this function in NumPy 2.0.

• Changes

  • Complex-valued {func}jax.numpy.geomspace now chooses the logarithmic spiral branch consistent with that of NumPy 2.0.
  • The behavior of lax.rng_bit_generator, and in turn the 'rbg' and 'unsafe_rbg' PRNG implementations, under jax.vmap has changed so that mapping over keys results in random generation only from the first key in the batch.
  • Docs now use jax.random.key for construction of PRNG key arrays rather than jax.random.PRNGKey.

• Deprecations & Removals

  • {func}jax.tree_map is deprecated; use jax.tree.map instead, or for backward compatibility with older JAX versions, use {func}jax.tree_util.tree_map.
  • {func}jax.clear_backends is deprecated as it does not necessarily do what its name suggests and can lead to unexpected consequences, e.g., it will not destroy existing backends and release corresponding owned resources. Use {func}jax.clear_caches if you only want to clean up compilation caches. For backward compatibility or you really need to switch/reinitialize the default backend, use {func}jax.extend.backend.clear_backends.
  • The jax.experimental.maps module and jax.experimental.maps.xmap are deprecated. Use jax.experimental.shard_map or jax.vmap with the spmd_axis_name argument for expressing SPMD device-parallel computations.
  • The jax.experimental.host_callback module is deprecated. Use instead the new JAX external callbacks. Added JAX_HOST_CALLBACK_LEGACY flag to assist in the transition to the new callbacks. See {jax-issue}#20385 for a discussion.
  • Passing arguments to {func}jax.numpy.array_equal and {func}jax.numpy.array_equiv that cannot be converted to a JAX array now results in an exception.
  • The deprecated flag jax_parallel_functions_output_gda has been removed. This flag was long deprecated and did nothing; its use was a no-op.
  • The previously-deprecated imports jax.interpreters.ad.config and jax.interpreters.ad.source_info_util have now been removed. Use jax.config and jax.extend.source_info_util instead.
  • JAX export does not support older serialization versions anymore. Version 9 has been supported since October 27th, 2023 and has become the default since February 1, 2024. See a description of the versions. This change could break clients that set a specific JAX serialization version lower than 9.

• New Features

  • Added CUDA Array Interface import support (requires jaxlib 0.4.24).
  • JAX arrays now support NumPy-style scalar boolean indexing, e.g. x[True] or x[False].
  • Added {mod}jax.tree module, with a more convenient interface for referencing functions in {mod}jax.tree_util.
  • {func}jax.tree.transpose (i.e. {func}jax.tree_util.tree_transpose) now accepts inner_treedef=None, in which case the inner treedef will be automatically inferred.

• Changes

  • Pallas now uses XLA instead of the Triton Python APIs to compile Triton kernels. You can revert to the old behavior by setting the JAX_TRITON_COMPILE_VIA_XLA environment variable to "0".
  • Several deprecated APIs in {mod}jax.interpreters.xla that were removed in v0.4.24 have been re-added in v0.4.25, including backend_specific_translations, translations, register_translation, xla_destructure, TranslationRule, TranslationContext, and XLAOp. These are still considered deprecated, and will be removed again in the future when better replacements are available. Refer to {jax-issue}#19816 for discussion.

• Deprecations & Removals

  • {func}jax.numpy.linalg.solve now shows a deprecation warning for batched 1D solves with b.ndim > 1. In the future these will be treated as batched 2D solves.
  • Conversion of a non-scalar array to a Python scalar now raises an error, regardless of the size of the array. Previously a deprecation warning was raised in the case of non-scalar arrays of size 1. This follows a similar deprecation in NumPy.
  • The previously deprecated configuration APIs have been removed following a standard 3 months deprecation cycle (see {ref}api-compatibility). These include
    • the jax.config.config object and
    • the define_*_state and DEFINE_* methods of {data}jax.config.
  • Importing the jax.config submodule via import jax.config is deprecated. To configure JAX use import jax and then reference the config object via jax.config.
  • The minimum jaxlib version is now 0.4.20.

• Changes

  • JAX lowering to StableHLO does not depend on physical devices anymore. If your primitive wraps custom_partitioning or JAX callbacks in the lowering rule i.e. function passed to rule parameter of mlir.register_lowering then add your primitive to jax._src.dispatch.prim_requires_devices_during_lowering set. This is needed because custom_partitioning and JAX callbacks need physical devices to create Shardings during lowering. This is a temporary state until we can create Shardings without physical devices.
  • {func}jax.numpy.argsort and {func}jax.numpy.sort now support the stable and descending arguments.
  • Several changes to the handling of shape polymorphism (used in {mod}jax.experimental.jax2tf and {mod}jax.experimental.export):
    • cleaner pretty-printing of symbolic expressions ({jax-issue}#19227)
    • added the ability to specify symbolic constraints on the dimension variables. This makes shape polymorphism more expressive, and gives a way to workaround limitations in the reasoning about inequalities. See https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#user-specified-symbolic-constraints.
    • with the addition of symbolic constraints ({jax-issue}#19235) we now consider dimension variables from different scopes to be different, even if they have the same name. Symbolic expressions from different scopes cannot interact, e.g., in arithmetic operations. Scopes are introduced by {func}jax.experimental.jax2tf.convert, {func}jax.experimental.export.symbolic_shape, {func}jax.experimental.export.symbolic_args_specs. The scope of a symbolic expression e can be read with e.scope and passed into the above functions to direct them to construct symbolic expressions in a given scope. See https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#user-specified-symbolic-constraints.
    • simplified and faster equality comparisons, where we consider two symbolic dimensions to be equal if the normalized form of their difference reduces to 0 ({jax-issue}#19231; note that this may result in user-visible behavior changes)
    • improved the error messages for inconclusive inequality comparisons ({jax-issue}#19235).
    • the core.non_negative_dim API (introduced recently) was deprecated and core.max_dim and core.min_dim were introduced ({jax-issue}#18953) to express max and min for symbolic dimensions. You can use core.max_dim(d, 0) instead of core.non_negative_dim(d).
    • the shape_poly.is_poly_dim is deprecated in favor of export.is_symbolic_dim ({jax-issue}#19282).
    • the export.args_specs is deprecated in favor of export.symbolic_args_specs ({jax-issue}#19283`).
    • the shape_poly.PolyShape and jax2tf.PolyShape are deprecated, use strings for polymorphic shapes specifications ({jax-issue}#19284).
    • JAX default native serialization version is now 9. This is relevant for {mod}jax.experimental.jax2tf and {mod}jax.experimental.export. See description of version numbers.
  • Refactored the API for jax.experimental.export. Instead of from jax.experimental.export import export you should use now from jax.experimental import export. The old way of importing will continue to work for a deprecation period of 3 months.
  • Added {func}jax.scipy.stats.sem.
  • {func}jax.numpy.unique with return_inverse = True returns inverse indices reshaped to the dimension of the input, following a similar change to {func}numpy.unique in NumPy 2.0.
  • {func}jax.numpy.sign now returns x / abs(x) for nonzero complex inputs. This is consistent with the behavior of {func}numpy.sign in NumPy version 2.0.
  • {func}jax.scipy.special.logsumexp with return_sign=True now uses the NumPy 2.0 convention for the complex sign, x / abs(x). This is consistent with the behavior of {func}scipy.special.logsumexp in SciPy v1.13.
  • JAX now supports the bool DLPack type for both import and export. Previously bool values could not be imported and were exported as integers.

• Deprecations & Removals

  • A number of previously deprecated functions have been removed, following a standard 3+ month deprecation cycle (see {ref}api-compatibility). This includes:
    • From {mod}jax.core: TracerArrayConversionError, TracerIntegerConversionError, UnexpectedTracerError, as_hashable_function, collections, dtypes, lu, map, namedtuple, partial, pp, ref, safe_zip, safe_map, source_info_util, total_ordering, traceback_util, tuple_delete, tuple_insert, and zip.
    • From {mod}jax.lax: dtypes, itertools, naryop, naryop_dtype_rule, standard_abstract_eval, standard_naryop, standard_primitive, standard_unop, unop, and unop_dtype_rule.
    • The jax.linear_util submodule and all its contents.
    • The jax.prng submodule and all its contents.
    • From {mod}jax.random: PRNGKeyArray, KeyArray, default_prng_impl, threefry_2x32, threefry2x32_key, threefry2x32_p, rbg_key, and unsafe_rbg_key.
    • From {mod}jax.tree_util: register_keypaths, AttributeKeyPathEntry, and GetItemKeyPathEntry.
    • from {mod}jax.interpreters.xla: backend_specific_translations, translations, register_translation, xla_destructure, TranslationRule, TranslationContext, axis_groups, ShapedArray, ConcreteArray, AxisEnv, backend_compile, and XLAOp.
    • from {mod}jax.numpy: NINF, NZERO, PZERO, row_stack, issubsctype, trapz, and in1d.
    • from {mod}jax.scipy.linalg: tril and triu.
  • The previously-deprecated method PRNGKeyArray.unsafe_raw_array has been removed. Use {func}jax.random.key_data instead.
  • bool(empty_array) now raises an error rather than returning False. This previously raised a deprecation warning, and follows a similar change in NumPy.
  • Support for the mhlo MLIR dialect has been deprecated. JAX no longer uses the mhlo dialect, in favor of stablehlo. APIs that refer to "mhlo" will be removed in the future. Use the "stablehlo" dialect instead.
  • {mod}jax.random: passing batched keys directly to random number generation functions, such as {func}~jax.random.bits, {func}~jax.random.gamma, and others, is deprecated and will emit a FutureWarning. Use jax.vmap for explicit batching.
  • {func}jax.lax.tie_in is deprecated: it has been a no-op since JAX v0.2.0.

• Fixed a bug that caused verbose logging from the GPU compiler during compilation.

• Deprecations
  • The device_buffer and device_buffers properties of JAX arrays are deprecated. Explicit buffers have been replaced by the more flexible array sharding interface, but the previous outputs can be recovered this way:
    • arr.device_buffer becomes arr.addressable_data(0)
    • arr.device_buffers becomes [x.data for x in arr.addressable_shards]

• New Features

  • Added {obj}jax.nn.squareplus.

• Changes

  • The minimum jaxlib version is now 0.4.19.
  • Released wheels are built now with clang instead of gcc.
  • Enforce that the device backend has not been initialized prior to calling jax.distributed.initialize().
  • Automate arguments to jax.distributed.initialize() in cloud TPU environments.

• Deprecations

  • The previously-deprecated sym_pos argument has been removed from {func}jax.scipy.linalg.solve. Use assume_a='pos' instead.
  • Passing None to {func}jax.array or {func}jax.asarray, either directly or within a list or tuple, is deprecated and now raises a {obj}FutureWarning. It currently is converted to NaN, and in the future will raise a {obj}TypeError.
  • Passing the condition, x, and y parameters to jax.numpy.where by keyword arguments has been deprecated, to match numpy.where.
  • Passing arguments to {func}jax.numpy.array_equal and {func}jax.numpy.array_equiv that cannot be converted to a JAX array is deprecated and now raises a {obj}DeprecationWaning. Currently the functions return False, in the future this will raise an exception.
  • The device() method of JAX arrays is deprecated. Depending on the context, it may be replaced with one of the following:
    • {meth}jax.Array.devices returns the set of all devices used by the array.
    • {attr}jax.Array.sharding gives the sharding configuration used by the array.

• Bug fixes
  • Fixed some type confusion between E4M3 and E5M2 float8 types.

• New Features

  • Added {obj}jax.typing.DTypeLike, which can be used to annotate objects that are convertible to JAX dtypes.
  • Added jax.numpy.fill_diagonal.

• Changes

  • JAX now requires SciPy 1.9 or newer.

• Bug fixes

  • Only process 0 in a multicontroller distributed JAX program will write persistent compilation cache entries. This fixes write contention if the cache is placed on a network file system such as GCS.
  • The version check for cusolver and cufft no longer considers the patch versions when determining if the installed version of these libraries is at least as new as the versions against which JAX was built.

• Changes

  • CUDA jaxlibs now depend on the user to install a compatible NCCL version. If using the recommended cuda12_pip installation, NCCL should be installed automatically. Currently, NCCL 2.16 or newer is required.
  • We now provide Linux aarch64 wheels, both with and without NVIDIA GPU support.
  • {meth}jax.Array.item now supports optional index arguments.

• Deprecations

  • A number of internal utilities and inadvertent exports in {mod}jax.lax have been deprecated, and will be removed in a future release.
    • jax.lax.dtypes: use jax.dtypes instead.
    • jax.lax.itertools: use itertools instead.
    • naryop, naryop_dtype_rule, standard_abstract_eval, standard_naryop, standard_primitive, standard_unop, unop, and unop_dtype_rule are internal utilities, now deprecated without replacement.

• Bug fixes

  • Fixed Cloud TPU regression where compilation would OOM due to smem.