search

facebookresearch / fmmax

Fourier modal method with Jax
MIT License
97 stars 10 forks source link

Update jax requirement from <=0.4.31 to <=0.4.34 #133

Closed dependabot[bot] closed 2 weeks ago

dependabot[bot] commented 3 weeks ago

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.4.34

  • New Functionality

    • This release includes wheels for Python 3.13. Free-threading mode is not yet supported.
    • jax.errors.JaxRuntimeError has been added as a public alias for the formerly private XlaRuntimeError type.

  • Breaking changes

    • jax_pmap_no_rank_reduction flag is set to True by default.
      • array[0] on a pmap result now introduces a reshape (use array[0:1] instead).
      • The per-shard shape (accessable via jax_array.addressable_shards or jax_array.addressable_data(0)) now has a leading (1, ...). Update code that directly accesses shards accordingly. The rank of the per-shard-shape now matches that of the global shape which is the same behavior as jit. This avoids costly reshapes when passing results from pmap into jit.
    • jax.experimental.host_callback has been deprecated since March 2024, with JAX version 0.4.26. Now we set the default value of the --jax_host_callback_legacy configuration value to True, which means that if your code uses jax.experimental.host_callback APIs, those API calls will be implemented in terms of the new jax.experimental.io_callback API. If this breaks your code, for a very limited time, you can set the --jax_host_callback_legacy to True. Soon we will remove that configuration option, so you should instead transition to using the new JAX callback APIs. See #20385 for a discussion.

  • Deprecations

    • In jax.numpy.trim_zeros, non-arraylike arguments or arraylike arguments with ndim != 1 are now deprecated, and in the future will result in an error.
    • Internal pretty-printing tools jax.core.pp_* have been removed, after being deprecated in JAX v0.4.30.
    • jax.lib.xla_client.Device is deprecated; use jax.Device instead.
    • jax.lib.xla_client.XlaRuntimeError has been deprecated. Use jax.errors.JaxRuntimeError instead.
    • The default behavior of jax.pure_callback and jax.extend.ffi.ffi_call under vmap has been deprecated and so has the vectorized parameter to those functions. The vmap_method parameter should be used instead for better defined behavior. See the discussion in #23881 for more details.

  • Deletion:

    • jax.xla_computation is deleted. It has been 3 months since its deprecation in 0.4.30 JAX release. Please use the AOT APIs to get the same functionality as jax.xla_computation.
      • jax.xla_computation(fn)(*args, **kwargs) can be replaced with jax.jit(fn).lower(*args, **kwargs).compiler_ir('hlo').
      • You can also use .out_info property of jax.stages.Lowered to get the output information (like tree structure, shape and dtype).
      • For cross-backend lowering, you can replace

... (truncated)

Changelog

Sourced from jax's changelog.

jax 0.4.34 (October 4, 2023)

  • New Functionality

    • This release includes wheels for Python 3.13. Free-threading mode is not yet supported.
    • jax.errors.JaxRuntimeError has been added as a public alias for the formerly private XlaRuntimeError type.

  • Breaking changes

    • jax_pmap_no_rank_reduction flag is set to True by default.
      • array[0] on a pmap result now introduces a reshape (use array[0:1] instead).
      • The per-shard shape (accessable via jax_array.addressable_shards or jax_array.addressable_data(0)) now has a leading (1, ...). Update code that directly accesses shards accordingly. The rank of the per-shard-shape now matches that of the global shape which is the same behavior as jit. This avoids costly reshapes when passing results from pmap into jit.
    • jax.experimental.host_callback has been deprecated since March 2024, with JAX version 0.4.26. Now we set the default value of the --jax_host_callback_legacy configuration value to True, which means that if your code uses jax.experimental.host_callback APIs, those API calls will be implemented in terms of the new jax.experimental.io_callback API. If this breaks your code, for a very limited time, you can set the --jax_host_callback_legacy to True. Soon we will remove that configuration option, so you should instead transition to using the new JAX callback APIs. See {jax-issue}[#20385](https://github.com/jax-ml/jax/issues/20385) for a discussion.

  • Deprecations

    • In {func}jax.numpy.trim_zeros, non-arraylike arguments or arraylike arguments with ndim != 1 are now deprecated, and in the future will result in an error.
    • Internal pretty-printing tools jax.core.pp_* have been removed, after being deprecated in JAX v0.4.30.
    • jax.lib.xla_client.Device is deprecated; use jax.Device instead.
    • jax.lib.xla_client.XlaRuntimeError has been deprecated. Use jax.errors.JaxRuntimeError instead.
    • The default behavior of {func}jax.pure_callback and {func}jax.extend.ffi.ffi_call under vmap has been deprecated and so has the vectorized parameter to those functions. The vmap_method parameter should be used instead for better defined behavior. See the discussion in {jax-issue}[#23881](https://github.com/jax-ml/jax/issues/23881) for more details.

  • Deletion:

    • jax.xla_computation is deleted. It's been 3 months since it's deprecation in 0.4.30 JAX release. Please use the AOT APIs to get the same functionality as jax.xla_computation.
      • jax.xla_computation(fn)(*args, **kwargs) can be replaced with jax.jit(fn).lower(*args, **kwargs).compiler_ir('hlo').
      • You can also use .out_info property of jax.stages.Lowered to get the output information (like tree structure, shape and dtype).

... (truncated)

Commits
  • affba36 JAX 0.4.34 release
  • 59cb692 Update XLA dependency to use revision
  • 8fc8082 Merge pull request #24077 from jakevdp:badge-links
  • cfb7541 Merge pull request #23738 from keshavb96:disable_remat_pass
  • 78b65dd Merge pull request #24070 from jakevdp:issubdtype-doc
  • a8744e3 README: fix badge href links
  • 152a873 Update XLA dependency to use revision
  • c9f946e Only thread a discharged ref value through a cond when it changes in some branch
  • 816947b Merge pull request #24071 from hawkinsp:winci
  • e212c77 Merge pull request #23891 from ROCm:build-fixes-rollup
  • 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 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[bot] commented 2 weeks ago

OK, I won't notify you again about this release, but will get in touch when a new version is available. If you'd rather skip all updates until the next major or minor version, let me know by commenting @dependabot ignore this major version or @dependabot ignore this minor version. You can also ignore all major, minor, or patch releases for a dependency by adding an ignore condition with the desired update_types to your config file.

If you change your mind, just re-open this PR and I'll resolve any conflicts on it.