Bump jax from 0.7.2 to 0.8.0

[dependabot]

Bumps jax from 0.7.2 to 0.8.0.

Changelog

Sourced from jax's changelog.

Change log

Best viewed here. For the changes specific to the experimental Pallas APIs, see {ref}pallas-changelog.

JAX follows Effort-based versioning; for a discussion of this and JAX's API compatibility policy, refer to {ref}api-compatibility. For the Python and NumPy version support policy, refer to {ref}version-support-policy.

Unreleased

• Breaking changes:

  • JAX is changing the default jax.pmap implementation to one implemented in terms of jax.jit and jax.shard_map. jax.pmap is in maintenance mode and we encourage all new code to use jax.shard_map directly. See the migration guide for more information.
  • The auto= parameter of jax.experimental.shard_map.shard_map has been removed. This means that jax.experimental.shard_map.shard_map no longer supports nesting. If you want to nest shard_map calls, please use jax.shard_map.
  • JAX no longer allows passing objects that support __jax_array__ directly to, e.g. jit-ed functions. Call jax.numpy.asarray on them first.
  • {func}jax.numpy.cov is now returns NaN for empty arrays ({jax-issue}[#32305](https://github.com/jax-ml/jax/issues/32305)), and matches NumPy 2.2 behavior for single-row design matrices ({jax-issue}[#32308](https://github.com/jax-ml/jax/issues/32308)).
  • JAX no longer accepts Array values where a dtype value is expected. Call .dtype on these values first.
  • The deprecated function {func}jax.interpreters.mlir.custom_call was removed.
  • The jax.util, jax.extend.ffi, and jax.experimental.host_callback modules have been removed. All public APIs within these modules were deprecated and removed in v0.7.0 or earlier.
  • The deprecated symbol {obj}jax.custom_derivatives.custom_jvp_call_jaxpr_p was removed.
  • jax.experimental.multihost_utils.process_allgather raises an error when the input is a jax.Array and not fully-addressable and tiled=False. To fix this, pass tiled=True to your process_allgather invocation.
  • from {mod}jax.experimental.compilation_cache, the deprecated symbols is_initialized and initialize_cache were removed.
  • The deprecated function {func}jax.interpreters.xla.canonicalize_dtype was removed.
  • {mod}jaxlib.hlo_helpers has been removed. Use {mod}jax.ffi instead.

... (truncated)

Commits

• See full diff 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 <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)