Bump jax from 0.8.0 to 0.8.1 in the patch group

[dependabot]

Bumps the patch group with 1 update: jax.

Updates jax from 0.8.0 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.

Changelog

Sourced from jax's changelog.

JAX 0.8.1 (November 18, 2025)

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

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

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually 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 <dependency name> major version will close this group update PR and stop Dependabot creating any more for the specific dependency's major version (unless you unignore this specific dependency's major version or upgrade to it yourself)
• @dependabot ignore <dependency name> minor version will close this group update PR and stop Dependabot creating any more for the specific dependency's minor version (unless you unignore this specific dependency's minor version or upgrade to it yourself)
• @dependabot ignore <dependency name> will close this group update PR and stop Dependabot creating any more for the specific dependency (unless you unignore this specific dependency or upgrade to it yourself)
• @dependabot unignore <dependency name> will remove all of the ignore conditions of the specified dependency
• @dependabot unignore <dependency name> <ignore condition> will remove the ignore condition of the specified dependency and ignore conditions