search

team-ocean / veros

The versatile ocean simulator, in pure Python, powered by JAX.
https://veros.readthedocs.io
MIT License
330 stars 55 forks source link

⬆️ Bump jax from 0.4.29 to 0.4.30 #615

Closed dependabot[bot] closed 4 months ago

dependabot[bot] commented 4 months ago

Bumps jax from 0.4.29 to 0.4.30.

Changelog

Sourced from jax's changelog.

jax 0.4.30 (June 18, 2024)

  • Changes

    • JAX supports ml_dtypes >= 0.2. In 0.4.29 release, the ml_dtypes version was bumped to 0.4.0 but this has been rolled back in this release to give users of both TensorFlow and JAX more time to migrate to a newer TensorFlow release.
    • jax.experimental.mesh_utils can now create an efficient mesh for TPU v5e.
    • jax now depends on jaxlib directly. This change was enabled by the CUDA plugin switch: there are no longer multiple jaxlib variants. You can install a CPU-only jax with pip install jax, no extras required.
    • Added an API for exporting and serializing JAX functions. This used to exist in jax.experimental.export (which is being deprecated), and will now live in jax.export. See the documentation.

  • Deprecations

    • Internal pretty-printing tools jax.core.pp_* are deprecated, and will be removed in a future release.
    • Hashing of tracers is deprecated, and will lead to a TypeError in a future JAX release. This previously was the case, but there was an inadvertent regression in the last several JAX releases.
    • jax.experimental.export is deprecated. Use {mod}jax.export instead. See the migration guide.
    • Passing an array in place of a dtype is now deprecated in most cases; e.g. for arrays x and y, x.astype(y) will raise a warning. To silence it use x.astype(y.dtype).
    • jax.xla_computation is deprecated and will be removed in a future 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 jax.xla_computation(fn, backend='tpu')(*args, **kwargs) with jax.jit(fn).trace(*args, **kwargs).lower(lowering_platforms=('tpu',)).compiler_ir('hlo').

jaxlib 0.4.30 (June 18, 2024)

  • Support for monolithic CUDA jaxlibs has been dropped. You must use the plugin-based installation (pip install jax[cuda12] or pip install jax[cuda12_local]).
Commits
  • f4158ac Merge pull request #21949 from hawkinsp:winwheel
  • b0b0268 Add --allow-downgrade to Windows wheel builds.
  • 5d35c99 Merge pull request #21945 from hawkinsp:release
  • d7bc6b4 Update XLA dependency to use revision
  • 3f4f79c Prepare for 0.4.30 release.
  • dfcfb36 Pallas GPU no longer falls back to lax.pow for integer powers
  • 5bfd6af Removed unnecessary skip in pallas_test.py::SoftmaxTest
  • 3fd9326 [jax] Enable api_test with XLA:CPU thunks
  • ed4958c [XLA:Mosaic] Add internal scratch VMEM
  • 701c63e [Pallas/TPU] Add API for megacore partitioning of pipelines
  • 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 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)