build(deps-dev): bump jax from 0.5.0 to 0.6.2

[dependabot]

Bumps jax from 0.5.0 to 0.6.2.

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.

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.

JAX v0.6.0

• Breaking changes

  • 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 [#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]

... (truncated)

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

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)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.15%. Comparing base (b9c6615) to head (c988731).
⚠️ Report is 14 commits behind head on devel.

Additional details and impacted files

@@            Coverage Diff             @@
##            devel    #5102      +/-   ##
==========================================
+ Coverage   82.14%   82.15%   +0.01%     
==========================================
  Files         709      709              
  Lines       72453    72468      +15     
  Branches     3615     3616       +1     
==========================================
+ Hits        59514    59536      +22     
+ Misses      11775    11769       -6     
+ Partials     1164     1163       -1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[njzjz]

Need TensorFlow to bump to a specific version as well. This is the limitation of jax2tf.

[njzjz]

Try 2.19.1