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

Note
Automatic rebases have been disabled on this pull request as it has been open for over 30 days.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.45%. Comparing base (eab3419) to head (f03491d).

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5102      +/-   ##
==========================================
- Coverage   82.47%   82.45%   -0.03%     
==========================================
  Files         825      825              
  Lines       87726    87727       +1     
  Branches     4206     4206              
==========================================
- Hits        72355    72335      -20     
- Misses      14088    14110      +22     
+ Partials     1283     1282       -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

[dependabot]

A newer version of jax exists, but since this PR has been edited by someone other than Dependabot I haven't updated it. You'll get a PR for the updated version as normal once this PR is merged.

[njzjz]

@dependabot recreate

[njzjz]

Possible reviewers based on changed lines, exact file history, and exact-file review history:

• @wanghan-iapcm — 1 commits on changed files; 18 reviews on exact changed files (pyproject.toml).

No review request was made automatically.

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh