Update jax requirement from ~=0.4.38 to ~=0.8.1 in /dev_tools/requirements/deps

[dependabot[bot]]

Updates the requirements on jax to permit the latest version.

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.

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

Most Recent Ignore Conditions Applied to This Pull Request

Dependency Name	Ignore Conditions
jax	[>= 0.5.dev0, < 0.6]

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)