# JAX v0.6.1 — JAX v0.6.1

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

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

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

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

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