chore(deps-dev): update jax requirement from <0.6,>=0.4 to >=0.4,<0.7

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.6.2

• New features:

  • Added jax.tree.broadcast which implements a pytree prefix broadcasting helper.

• Changes

  • The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12.

Changelog

Sourced from jax's changelog.

JAX 0.6.2 (June 17, 2025)

• New features:

  • Added {func}jax.tree.broadcast which implements a pytree prefix broadcasting helper.

• Changes

  • The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12.

JAX 0.6.1 (May 21, 2025)

• New features:

  • Added {func}jax.lax.axis_size which returns the size of the mapped axis given its name.

• Changes

  • Additional checking for the versions of CUDA package dependencies was re-enabled, having been accidentally disabled in a previous release.
  • JAX nightly packages are now published to artifact registry. To install these packages, see the JAX installation guide.
  • jax.sharding.PartitionSpec no longer inherits from a tuple.
  • jax.ShapeDtypeStruct is immutable now. Please use .update method to update your ShapeDtypeStruct instead of doing in-place updates.

• Deprecations

  • jax.custom_derivatives.custom_jvp_call_jaxpr_p is deprecated, and will be removed in JAX v0.7.0.

JAX 0.6.0 (April 16, 2025)

• Breaking changes

  • {func}jax.numpy.array no longer accepts None. This behavior was deprecated since November 2023 and is now removed.
  • Removed the config.jax_data_dependent_tracing_fallback config option, which was added temporarily in v0.4.36 to allow users to opt out of the new "stackless" tracing machinery.
  • Removed the config.jax_eager_pmap config option.
  • Disallow the calling of lower and trace AOT APIs on the result of jax.jit if there have been subsequent wrappers applied. Previously this worked, but silently ignored the wrappers. The workaround is to apply jax.jit last among the wrappers, and similarly for jax.pmap. See {jax-issue}[#27873](https://github.com/jax-ml/jax/issues/27873).
  • The cuda12_pip extra for jax has been removed; use pip install jax[cuda12] instead.

• Changes

  • The minimum CuDNN version is v9.8.
  • JAX is now built using CUDA 12.8. All versions of CUDA 12.1 or newer remain supported.

... (truncated)

Commits

• 1ad05bb Add LLVM patch to fix AVX512 codegeneration problem
• 8f81490 Prepare for JAX release 0.6.2
• e4de90e Update XLA dependency to use revision
• 02688e1 [Pallas][Mosaic GPU] Enable collective MMA from TMEM.
• dc9ef61 Merge pull request #29410 from DanisNone:nn-type
• 353e7fa Merge pull request #29516 from jakevdp:enable-x64-warning
• 3d37b0d Merge pull request #29504 from vfdev-5:tsan-ft-removed-fixed-suppression
• c2cc9f9 [pallas] AbstractMemoryRef now implements all functional update methods via...
• 0fd0821 [Pallas TPU] Add flag to enable using registers to keep track of slot info
• f22896a jax.experimental.enable_x64: add warning to docstring
• Additional commits viewable in compare view