Bump jax from 0.8.0 to 0.8.1

[dependabot]

Bumps 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 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)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.98%. Comparing base (8ac3ec3) to head (eb94dbe).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #623   +/-   ##
=======================================
  Coverage   96.98%   96.98%           
=======================================
  Files          42       42           
  Lines        1094     1094           
=======================================
  Hits         1061     1061           
  Misses         33       33           

Flag	Coverage Δ
unittests	96.98% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Thank you for your contribution @dependabot[bot] 🚀! Your github-pages is ready for download 👉 here 👈!
_{(The artifact expires on 2025-12-01T05:33:12Z. You can re-generate it by re-running the workflow here.)}