# JAX v0.5.1 — JAX v0.5.1

- Product: JAX (https://whatsnew.fyi/product/jax)
- Vendor: Google
- 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

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

---

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