Update jax requirement from <0.6.1,>=0.4.25 to >=0.4.25,<0.7.3

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.7.2

• Breaking changes:

  • jax.dlpack.from_dlpack no longer accepts a DLPack capsule. This behavior was deprecated and is now removed. The function must be called with an array implementing __dlpack__ and __dlpack_device__.

• Changes

  • The minimum supported NumPy version is now 2.0. Since SciPy 1.13 is required for NumPy 2.0 support, the minimum supported SciPy version is now 1.13.

  • JAX now represents constants in its internal jaxpr representation as a LiteralArray, which is a private JAX type that duck types as a numpy.ndarray. This type may be exposed to users via custom_jvp rules, for example, and may break code that uses isinstance(x, np.ndarray). If this breaks your code, you may convert these arrays to classic NumPy arrays using np.asarray(x).

• Bug fixes

  • arr.view(dtype=None) now returns the array unchanged, matching NumPy's semantics. Previously it returned the array with a float dtype.
  • jax.random.randint now produces a less-biased distribution for 8-bit and 16-bit integer types ({jax-issue}[#27742](https://github.com/jax-ml/jax/issues/27742)). To restore the previous biased behavior, you may temporarily set the jax_safer_randint configuration to False, but note this is a temporary config that will be removed in a future release.

• Deprecations:

  • The parameters enable_xla and native_serialization for jax2tf.convert are deprecated and will be removed in a future version of JAX. These were used for jax2tf with non-native serialization, which has been now removed.
  • Setting the config state jax_pmap_no_rank_reduction to False is deprecated. By default, jax_pmap_no_rank_reduction will be set to True and jax.pmap shards will not have their rank reduced, keeping the same rank as their enclosing array.

Changelog

Sourced from jax's changelog.

JAX 0.7.2 (September 16, 2025)

• Breaking changes:

  • {func}jax.dlpack.from_dlpack no longer accepts a DLPack capsule. This behavior was deprecated and is now removed. The function must be called with an array implementing __dlpack__ and __dlpack_device__.

• Changes

  • The minimum supported NumPy version is now 2.0. Since SciPy 1.13 is required for NumPy 2.0 support, the minimum supported SciPy version is now 1.13.

  • JAX now represents constants in its internal jaxpr representation as a TypedNdArray, which is a private JAX type that duck types as a numpy.ndarray. This type may be exposed to users via custom_jvp rules, for example, and may break code that uses isinstance(x, np.ndarray). If this breaks your code, you may convert these arrays to classic NumPy arrays using np.asarray(x).

• Bug fixes

  • arr.view(dtype=None) now returns the array unchanged, matching NumPy's semantics. Previously it returned the array with a float dtype.
  • jax.random.randint now produces a less-biased distribution for 8-bit and 16-bit integer types ({jax-issue}[#27742](https://github.com/jax-ml/jax/issues/27742)). To restore the previous biased behavior, you may temporarily set the jax_safer_randint configuration to False, but note this is a temporary config that will be removed in a future release.

• Deprecations:

  • The parameters enable_xla and native_serialization for jax2tf.convert are deprecated and will be removed in a future version of JAX. These were used for jax2tf with non-native serialization, which has been now removed.
  • Setting the config state jax_pmap_no_rank_reduction to False is deprecated. By default, jax_pmap_no_rank_reduction will be set to True and jax.pmap shards will not have their rank reduced, keeping the same rank as their enclosing array.

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

... (truncated)

Commits

• 9423314 Prepare for JAX release 0.7.2
• 6093c04 Fix device_put's abstract_eval to return the correct sharding on the type whe...
• a29751e [Pallas][Mosaic GPU] Improve ragged dot kernel for blackwell.
• 8084f44 Merge pull request #31790 from mattjj:hijax-refs
• 1f75be0 [hijax] support refs to hijax types
• 9a238f1 [Pallas/Hijax] Let MemoryRef be parameterized by a type, instead of forcing i...
• a27ef01 [Mosaic] Add tpu.barrier.
• e86b6ae [Pallas:TPU] Use the lowering backend to query the libTPU version
• 2c21ce5 Merge pull request #31822 from mattjj:named-scope-docs-fix
• 920eb5e [Pallas:MGPU] Enable a test for WG semantic and clarify skip reasons.
• 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 #122.