Bump jax from 0.7.2 to 0.8.1

[dependabot]

Bumps jax from 0.7.2 to 0.8.1.

Release notes

Sourced from jax's releases.

JAX v0.8.1

• New features:

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

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

JAX v0.8.0

• 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

... (truncated)

Changelog

Sourced from jax's changelog.

JAX 0.8.1

• 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](https://github.com/jax-ml/jax/issues/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.

JAX 0.8.0 (October 15, 2025)

• Breaking changes:

... (truncated)

Commits

• 8fc84ad Prepare for JAX release 0.8.1
• d8e9e28 Add a TPU platform check in make_mesh in the slice_index check because GPUs a...
• 10fe299 Correct the SHA for the XLA archive at the current commit.
• 40024b1 Expose jax.sharding.get_mesh() as a way to get the concrete mesh. This is u...
• f8f68d6 Merge pull request #33374 from jakevdp:fix-lint
• 79e986b Fix lint issue at head
• e24d3ca Update XLA dependency to use revision
• 30ebb4f Merge pull request #33316 from jakevdp:doc-coverage-test
• cd11231 Add test of documentation coverage
• ad960fd [Mosaic] Allow padding in small tiling row shuffle reshape.
• Additional commits viewable in compare view

You can trigger a rebase of this PR by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)

Note
Automatic rebases have been disabled on this pull request as it has been open for over 30 days.

[dependabot]

A newer version of jax exists, but since this PR has been edited by someone other than Dependabot I haven't updated it. You'll get a PR for the updated version as normal once this PR is merged.

[dionhaefner]

@dependabot rebase

[dependabot]

Looks like this PR has been edited by someone other than Dependabot. That means Dependabot can't rebase it - sorry!

If you're happy for Dependabot to recreate it from scratch, overwriting any edits, you can request @dependabot recreate.

[dionhaefner]

@dependabot recreate