Update dependency jax to v0.7.0

[renovate]

This PR contains the following updates:

Package	Change	Age	Adoption	Passing	Confidence
jax	==0.6.1 -> ==0.7.0

---

Release Notes

jax-ml/jax (jax)

v0.7.0

Compare Source

• New features:

  • Added jax.P which is an alias for jax.sharding.PartitionSpec.
  • Added {func}jax.tree.reduce_associative.

• Breaking changes:

  • JAX is migrating from GSPMD to Shardy by default. See the
    migration guide
    for more information.
  • JAX autodiff is switching to using direct linearization by default (instead of
    implementing linearization via JVP and partial eval).
    See migration guide
    for more information.
  • jax.stages.OutInfo has been replaced with jax.ShapeDtypeStruct.
  • {func}jax.jit now requires fun to be passed by position, and additional
    arguments to be passed by keyword. Doing otherwise will result in an error
    starting in v0.7.x. This raised a DeprecationWarning in v0.6.x.
  • The minimum Python version is now 3.11. 3.11 will remain the minimum
    supported version until July 2026.
  • Layout API renames:
    • Layout, .layout, .input_layouts and .output_layouts have been
      renamed to Format, .format, .input_formats and .output_formats
    • DeviceLocalLayout, .device_local_layout have been renamed to Layout
      and .layout
  • jax.experimental.shard module has been deleted and all the APIs have been
    moved to the jax.sharding endpoint. So use jax.sharding.reshard,
    jax.sharding.auto_axes and jax.sharding.explicit_axes instead of their
    experimental endpoints.
  • lax.infeed and lax.outfeed were removed, after being deprecated in
    JAX 0.6. The transfer_to_infeed and transfer_from_outfeed methods were
    also removed the Device objects.
  • The jax.extend.core.primitives.pjit_p primitive has been renamed to
    jit_p, and its name attribute has changed from "pjit" to "jit".
    This affects the string representations of jaxprs. The same primitive is no
    longer exported from the jax.experimental.pjit module.
  • The (undocumented) function jax.extend.backend.add_clear_backends_callback
    has been removed. Users should use jax.extend.backend.register_backend_cache
    instead.

• Deprecations:

  • {obj}jax.dlpack.SUPPORTED_DTYPES is deprecated; please use the new
    {func}jax.dlpack.is_supported_dtype function.
  • {func}jax.scipy.special.sph_harm has been deprecated following a similar
    deprecation in SciPy; use {func}jax.scipy.special.sph_harm_y instead.
  • From {mod}jax.interpreters.xla, the previously deprecated symbols
    abstractify and pytype_aval_mappings have been removed.
  • {func}jax.interpreters.xla.canonicalize_dtype is deprecated. For
    canonicalizing dtypes, prefer {func}jax.dtypes.canonicalize_dtype.
    For checking whether an object is a valid jax input, prefer
    {func}jax.core.valid_jaxtype.
  • From {mod}jax.core, the previously deprecated symbols AxisName,
    ConcretizationTypeError, axis_frame, call_p, closed_call_p,
    get_type, trace_state_clean, typematch, and typecheck have been
    removed.
  • From {mod}jax.lib.xla_client, the previously deprecated symbols
    DeviceAssignment, get_topology_for_devices, and mlir_api_version
    have been removed.
  • jax.extend.ffi was removed after being deprecated in v0.5.0.
    Use {mod}jax.ffi instead.
  • {func}jax.lib.xla_bridge.get_compile_options is deprecated, and replaced by
    {func}jax.extend.backend.get_compile_options.

v0.6.2

Compare Source

• New features:

  • Added {func}jax.tree.broadcast which implements a pytree prefix broadcasting helper.

• Changes

  • The minimum NumPy version is 1.26 and the minimum SciPy version is 1.12.

---

Configuration

📅 Schedule: Branch creation - At any time (no schedule defined), Automerge - At any time (no schedule defined).

🚦 Automerge: Disabled by config. Please merge this manually once you are satisfied.

♻ Rebasing: Whenever PR becomes conflicted, or you tick the rebase/retry checkbox.

🔕 Ignore: Close this PR and you won't be reminded about this update again.

---

• If you want to rebase/retry this PR, check this box

---

This PR was generated by Mend Renovate. View the repository job log.

[coderabbitai]

Important

Review skipped

Bot user detected.

To trigger a single review, invoke the @coderabbitai review command.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

---

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Join our Discord community for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.