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

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

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

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

* Breaking changes * `jax.numpy.array` no longer accepts `None`. This behavior was deprecated since November 2023 and is now removed. * Removed the `config.jax_data_dependent_tracing_fallback` config option, which was added temporarily in v0.4.36 to allow users to opt out of the new "stackless" tracing machinery. * Removed the `config.jax_eager_pmap` config option. * Disallow the calling of `lower` and `trace` AOT APIs on the result of `jax.jit` if there have been subsequent wrappers applied. Previously this worked, but silently ignored the wrappers. The workaround is to apply `jax.jit` last among the wrappers, and similarly for `jax.pmap`. See `#27873`. * The `cuda12_pip` extra for `jax` has been removed; use `pip install jax[cuda12]` instead. * Changes * The minimum CuDNN version is v9.8. * JAX is now built using CUDA 12.8. All versions of CUDA 12.1 or newer remain supported. * JAX package extras are now updated to use dash instead of underscore to align with PEP 685. For instance, if you were previously using `pip install jax[cuda12_local]` to install JAX, run `pip install jax[cuda12-local]` instead. * `jax.jit` now requires `fun` to be passed by position, and additional arguments to be passed by keyword. Doing otherwise will result in a DeprecationWarning in v0.6.X, and an error in starting in v0.7.X. * Deprecations * `jax.tree_util.build_tree` is deprecated. Use `jax.tree.unflatten` instead. * Implemented host callback handlers for CPU and GPU devices using XLA's FFI and removed existing CPU/GPU handlers using XLA's custom call. * All APIs in `jax.lib.xla_extension` are now deprecated. * `jax.interpreters.mlir.hlo` and `jax.interpreters.mlir.func_dialect`, which were accidental exports, have been removed. If needed, they are available from `jax.extend.mlir`. * `jax.interpreters.mlir.custom_call` is deprecated. The APIs provided by `jax.ffi` should be used instead. * The deprecated use of `jax.ffi.ffi_call` with inline arguments is no longer supported. `jax.ffi.ffi_call` now unconditionally returns a callable. * The following exports in `jax.lib.xla_client` are deprecated: `get_topology_for_devices`, `heap_profile`, `mlir_api_version`, `Client`, `CompileOptions`, `DeviceAssignment`, `Frame`, `HloSharding`, `OpSharding`, `Traceback`. * The following internal APIs in `jax.util` are deprecated: `HashableFunction`, `as_hashable_function`, `cache`, `safe_map`, `safe_zip`, `split_dict`, `split_list`, `split_list_checked`, `split_merge`, `subvals`, `toposort`, `unzip2`, `wrap_name`, and `wraps`. * `jax.dlpack.to_dlpack` has been deprecated. You can usually pass a JAX `Array` directly to the `from_dlpack` function of another framework. If you need the functionality of `to_dlpack`, use the `__dlpack__` attribute of an array. * `jax.lax.infeed`, `jax.lax.infeed_p`, `jax.lax.outfeed`, and `jax.lax.outfeed_p` are deprecated and will be removed in JAX v0.7.0. * Several previously-deprecated APIs have been removed, including: * From `jax.lib.xla_client`: `ArrayImpl`, `FftType`, `PaddingType`, `PrimitiveType`, `XlaBuilder`, `dtype_to_etype`, `ops`, `register_custom_call_target`, `shape_from_pyval`, `Shape`, `XlaComputation`. * From `jax.lib.xla_extension`: `ArrayImpl`, `XlaRuntimeError`. * From `jax`: `jax.treedef_is_leaf`, `jax.tree_flatten`, `jax.tree_map`, `jax.tree_leaves`, `jax.tree_structure`, `jax.tree_transpose`, and `jax.tree_unflatten`. Replacements can be found in `jax.tree` or `jax.tree_util`. * From `jax.core`: `AxisSize`, `ClosedJaxpr`, `EvalTrace`, `InDBIdx`, `InputType`, `Jaxpr`, `JaxprEqn`, `Literal`, `MapPrimitive`, `OpaqueTraceState`, `OutDBIdx`, `Primitive`, `Token`, `TRACER_LEAK_DEBUGGER_WARNING`, `Var`, `concrete_aval`, `dedup_referents`, `escaped_tracer_error`, `extend_axis_env_nd`, `full_lower`, `get_referent`, `jaxpr_as_fun`, `join_effects`, `lattice_join`, `leaked_tracer_error`, `maybe_find_leaked_tracers`, `raise_to_shaped`, `raise_to_shaped_mappings`, `reset_trace_state`, `str_eqn_compact`, `substitute_vars_in_output_ty`, `typecompat`, and `used_axis_names_jaxpr`. Most have no public replacement, though a few are available at `jax.extend.core`. * The `vectorized` argument to `jax.pure_callback` and `jax.ffi.ffi_call`. Use the `vmap_method` parameter instead.

## jax 0.5.3 * 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.

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

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