# JAX v0.8.1 — JAX v0.8.1

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

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

* New features:

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

* Changes:

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

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

* Bug fixes:

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

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