Skip to content

⬆️ Bump jax[cuda12] from 0.4.38 to 0.6.1 #754

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: master
Choose a base branch
from

Conversation

dependabot[bot]
Copy link
Contributor

@dependabot dependabot bot commented on behalf of github Jun 2, 2025

Bumps jax[cuda12] from 0.4.38 to 0.6.1.

Release notes

Sourced from jax[cuda12]'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.

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] to install JAX, run pip install jax[cuda12-local] instead.
    • 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)

Changelog

Sourced from jax[cuda12]'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 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 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 compatibility score

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)

Bumps [jax[cuda12]](https://github.com/jax-ml/jax) from 0.4.38 to 0.6.1.
- [Release notes](https://github.com/jax-ml/jax/releases)
- [Changelog](https://github.com/jax-ml/jax/blob/main/CHANGELOG.md)
- [Commits](jax-ml/jax@jax-v0.4.38...jax-v0.6.1)

---
updated-dependencies:
- dependency-name: jax[cuda12]
  dependency-version: 0.6.1
  dependency-type: direct:development
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
@dependabot dependabot bot added dependencies Pull requests that update a dependency file python Pull requests that update Python code labels Jun 2, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
dependencies Pull requests that update a dependency file python Pull requests that update Python code
Projects
None yet
Development

Successfully merging this pull request may close these issues.

0 participants