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

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.6.1

• New features:

  • Added 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 reenabled, 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.

Changelog

Sourced from jax's changelog.

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.
  • JAX package extras are now updated to use dash instead of underscore to align with PEP 685. For instance, if you were previously using pip install jax[cuda12_local] to install JAX, run pip install jax[cuda12-local] instead.
  • {func}jax.jit now requires fun to be passed by position, and additional arguments to be passed by keyword. Doing otherwise will result in a DeprecationWarning in v0.6.X, and an error in starting in v0.7.X.

• Deprecations

... (truncated)

Commits

• 382506f Prepare for jax release v0.6.1
• 5f764b5 [Pallas/Mosaic GPU] Loosen tiling requirements for get and swap.
• 8653a78 [Mosaic GPU] Support s8xs8->s32 WGMMA.
• 2f0ca79 [pallas:mosaic_gpu] Removed GPU* aliases
• 482b812 Update XLA dependency to use revision
• 0f84433 Disable tests for the paged attention kernel
• 4d6e39e Disable the newly added tfrt targets that never worked
• d8804aa Add missing test skips (for too old jaxlib) + bump minimum libtpu version
• eef1f6c Support passing PartitionSpecs to ShapeDtypeStruct when there is a mesh in co...
• 56d4cb0 Merge pull request #28871 from hawkinsp:ci
• 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 #105.