Skip to content
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

Bump jax from 0.5.0 to 0.5.1 #925

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

Conversation

dependabot[bot]
Copy link
Contributor

@dependabot dependabot bot commented on behalf of github Feb 25, 2025

Bumps jax from 0.5.0 to 0.5.1.

Release notes

Sourced from jax's releases.

JAX v0.5.1

  • New Features

    • Added an experimental jax.experimental.custom_dce.custom_dce decorator to support customizing the behavior of opaque functions under JAX-level dead code elimination (DCE). See [#25956](https://github.com/jax-ml/jax/issues/25956) for more details.
    • Added low-level reduction APIs in {mod}jax.lax: jax.lax.reduce_sum, jax.lax.reduce_prod, jax.lax.reduce_max, jax.lax.reduce_min, jax.lax.reduce_and, jax.lax.reduce_or, and jax.lax.reduce_xor.
    • jax.lax.linalg.qr, and jax.scipy.linalg.qr, now support column-pivoting on CPU and GPU. See #20282 and #25955 for more details.
  • Changes

    • JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES now work as env vars. Before they could only be specified via jax.config or flags.
    • JAX_CPU_COLLECTIVES_IMPLEMENTATION now defaults to 'gloo', meaning multi-process CPU communication works out-of-the-box.
    • The jax[tpu] TPU extra no longer depends on the libtpu-nightly package. This package may safely be removed if it is present on your machine; JAX now uses libtpu instead.
  • Deprecations

    • The internal function linear_util.wrap_init and the constructor core.Jaxpr now must take a non-empty core.DebugInfo kwarg. For a limited time, a DeprecationWarning is printed if jax.extend.linear_util.wrap_init is used without debugging info. A downstream effect of this several other internal functions need debug info. This change does not affect public APIs. See jax-ml/jax#26480 for more detail.
  • Bug fixes

    • TPU runtime startup and shutdown time should be significantly improved on TPU v5e and newer (from around 17s to around 8s). If not already set, you may need to enable transparent hugepages in your VM image (sudo sh -c 'echo always > /sys/kernel/mm/transparent_hugepage/enabled'). We hope to improve this further in future releases.
    • Persistent compilation cache no longer writes access time file if JAX_COMPILATION_CACHE_MAX_SIZE is unset or set to -1, i.e. if the LRU eviction policy isn't enabled. This should improve performance when using the cache with large-scale network storage.
Changelog

Sourced from jax's changelog.

jax 0.5.1 (Feb 24, 2025)

  • New Features

    • Added an experimental {func}jax.experimental.custom_dce.custom_dce decorator to support customizing the behavior of opaque functions under JAX-level dead code elimination (DCE). See {jax-issue}[#25956](https://github.com/jax-ml/jax/issues/25956) for more details.
    • Added low-level reduction APIs in {mod}jax.lax: {func}jax.lax.reduce_sum, {func}jax.lax.reduce_prod, {func}jax.lax.reduce_max, {func}jax.lax.reduce_min, {func}jax.lax.reduce_and, {func}jax.lax.reduce_or, and {func}jax.lax.reduce_xor.
    • {func}jax.lax.linalg.qr, and {func}jax.scipy.linalg.qr, now support column-pivoting on CPU and GPU. See {jax-issue}[#20282](https://github.com/jax-ml/jax/issues/20282) and {jax-issue}[#25955](https://github.com/jax-ml/jax/issues/25955) for more details.
  • Changes

    • JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES now work as env vars. Before they could only be specified via jax.config or flags.
    • JAX_CPU_COLLECTIVES_IMPLEMENTATION now defaults to 'gloo', meaning multi-process CPU communication works out-of-the-box.
    • The jax[tpu] TPU extra no longer depends on the libtpu-nightly package. This package may safely be removed if it is present on your machine; JAX now uses libtpu instead.
  • Deprecations

    • The internal function linear_util.wrap_init and the constructor core.Jaxpr now must take a non-empty core.DebugInfo kwarg. For a limited time, a DeprecationWarning is printed if jax.extend.linear_util.wrap_init is used without debugging info. A downstream effect of this several other internal functions need debug info. This change does not affect public APIs. See jax-ml/jax#26480 for more detail.
  • Bug fixes

    • TPU runtime startup and shutdown time should be significantly improved on TPU v5e and newer (from around 17s to around 8s). If not already set, you may need to enable transparent hugepages in your VM image (sudo sh -c 'echo always > /sys/kernel/mm/transparent_hugepage/enabled'). We hope to improve this further in future releases.
    • Persistent compilation cache no longer writes access time file if JAX_COMPILATION_CACHE_MAX_SIZE is unset or set to -1, i.e. if the LRU eviction policy isn't enabled. This should improve performance when using the cache with large-scale network storage.
Commits
  • 07440f4 Prepare for JAX release 0.5.1
  • b0cfcb8 Update XLA dependency to use revision
  • 7d3c63e [sharding_in_types] Add more reshape sharding support
  • 908ff49 [mosaic_gpu] Warmup the kernel when doing CUPTI based profiling
  • 74b2e02 [pallas:mosaic_gpu] Use {min,max}imumf instead of {min,max}numf
  • e0e5457 Update XLA dependency to use revision
  • d695aa4 [sharding_in_types] Add sharding rules for the following primitives:
  • a75bea5 Update XLA dependency to use revision
  • 7c4fe2a [sharding_in_types] Allow auto_axes and explicit_axes to take numpy array...
  • 3407785 Reverts 6e83de5909b5dad6827c1682726c840c9688b32d
  • 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](https://github.com/jax-ml/jax) from 0.5.0 to 0.5.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.5.0...jax-v0.5.1)

---
updated-dependencies:
- dependency-name: jax
  dependency-type: direct:development
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <[email protected]>
@dependabot dependabot bot added dependencies Pull requests that update a dependency file python Pull requests that update Python code labels Feb 25, 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