Bump jax from 0.7.0 to 0.7.1

[dependabot]

Bumps jax from 0.7.0 to 0.7.1.

Release notes

Sourced from jax's releases.

JAX v0.7.1

• New features

  • JAX now ships Python 3.14 and 3.14t wheels.
  • JAX now ships Python 3.13t and 3.14t wheels on Mac. Previously we only offered free-threading builds on Linux.

• Changes

  • Exposed jax.set_mesh which acts as a global setter and a context manager. Removed jax.sharding.use_mesh in favor of jax.set_mesh.
  • JAX is now built using CUDA 12.9. All versions of CUDA 12.1 or newer remain supported.
  • jax.lax.dot now implements the general dot product via the optional dimension_numbers argument.

• Deprecations:

  • jax.lax.zeros_like_array is deprecated. Please use jax.numpy.zeros_like instead.
  • Attempting to import jax.experimental.host_callback now results in a DeprecationWarning, and will result in an ImportError starting in JAX v0.8.0. Its APIs have raised NotImplementedError since JAX version 0.4.35.
  • In jax.lax.dot, passing the precision and preferred_element_type arguments by position is deprecated. Pass them by explicit keyword instead.
  • Several dozen internal APIs have been deprecated from jax.interpreters.ad, jax.interpreters.batching, and jax.interpreters.partial_eval; they are used rarely if ever outside JAX itself, and most are deprecated without any public replacement.

Changelog

Sourced from jax's changelog.

JAX 0.7.1 (August 20, 2025)

• New features

  • JAX now ships Python 3.14 and 3.14t wheels.
  • JAX now ships Python 3.13t and 3.14t wheels on Mac. Previously we only offered free-threading builds on Linux.

• Changes

  • Exposed jax.set_mesh which acts as a global setter and a context manager. Removed jax.sharding.use_mesh in favor of jax.set_mesh.
  • JAX is now built using CUDA 12.9. All versions of CUDA 12.1 or newer remain supported.
  • {func}jax.lax.dot now implements the general dot product via the optional dimension_numbers argument.

• Deprecations:

  • {func}jax.lax.zeros_like_array is deprecated. Please use {func}jax.numpy.zeros_like instead.
  • Attempting to import {mod}jax.experimental.host_callback now results in a DeprecationWarning, and will result in an ImportError starting in JAX v0.8.0. Its APIs have raised NotImplementedError since JAX version 0.4.35.
  • In {func}jax.lax.dot, passing the precision and preferred_element_type arguments by position is deprecated. Pass them by explicit keyword instead.
  • Several dozen internal APIs have been deprecated from {mod}jax.interpreters.ad, {mod}jax.interpreters.batching, and {mod}jax.interpreters.partial_eval; they are used rarely if ever outside JAX itself, and most are deprecated without any public replacement.

Commits

• 5712de4 Prepare for JAX release 0.7.1
• 6248205 [Mosaic GPU] Separate hardware-agnostic from hardware-specific passes in BUIL...
• a44869e [Mosaic GPU] Annotate alignment of GMEM pointers to be 256
• 1637155 [Mosaic GPU] Second attempt at fixing test_tcgen05_collective_mma.
• ffc0afb [Mosaic GPU] Attempt no. 2 at fixing the Mac OS build.
• 7ad1508 [Mosaic GPU] Do not predicate collective ops in the lowering.
• cbd0844 [Mosaic GPU] Fold non-gather indices into the gather TMA column index
• 89628e4 Update XLA dependency to use revision http://github.com/openxla/xla/commit/31...
• b09cc05 [Mosaic GPU] Include unistd.h directly in order to fix build on Mac.
• e44e8e6 Merge pull request #31151 from mattjj:fix-error-message
• 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)

[dependabot]

Superseded by #753.