chore(deps-dev): update jax requirement from <=0.2.12 to <=0.5.0

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.5.0

As of this release, JAX now uses effort-based versioning. Since this release makes a breaking change to PRNG key semantics that may require users to update their code, we are bumping the "meso" version of JAX to signify this.

• Breaking changes

  • Enable jax_threefry_partitionable by default (see the update note).

  • This release drops support for Mac x86 wheels. Mac ARM of course remains supported. For a recent discussion, see Will macOS x86 still be supported for the next few years? jax-ml/jax#22936.

    Two key factors motivated this decision:

    • The Mac x86 build (only) has a number of test failures and crashes. We would prefer to ship no release than a broken release.
    • Mac x86 hardware is end-of-life and cannot be easily obtained for developers at this point. So it is difficult for us to fix this kind of problem even if we wanted to.

    We are open to readding support for Mac x86 if the community is willing to help support that platform: in particular, we would need the JAX test suite to pass cleanly on Mac x86 before we could ship releases again.

• Changes:

  • The minimum NumPy version is now 1.25. NumPy 1.25 will remain the minimum supported version until June 2025.
  • The minimum SciPy version is now 1.11. SciPy 1.11 will remain the minimum supported version until June 2025.
  • jax.numpy.einsum now defaults to optimize='auto' rather than optimize='optimal'. This avoids exponentially-scaling trace-time in the case of many arguments ([#25214](https://github.com/jax-ml/jax/issues/25214)).
  • jax.numpy.linalg.solve no longer supports batched 1D arguments on the right hand side. To recover the previous behavior in these cases, use solve(a, b[..., None]).squeeze(-1).

• New Features

  • jax.numpy.fft.fftn, jax.numpy.fft.rfftn, jax.numpy.fft.ifftn, and jax.numpy.fft.irfftn now support transforms in more than 3 dimensions, which was previously the limit. See [#25606](https://github.com/jax-ml/jax/issues/25606) for more details.
  • Support added for user defined state in the FFI via the new jax.ffi.register_ffi_type_id function.
  • The AOT lowering .as_text() method now supports the debug_info option to include debugging information, e.g., source location, in the output.

• Deprecations

  • From jax.interpreters.xla, abstractify and pytype_aval_mappings are now deprecated, having been replaced by symbols of the same name in jax.core.

... (truncated)

Changelog

Sourced from jax's changelog.

jax 0.5.0 (Jan 17, 2025)

As of this release, JAX now uses effort-based versioning. Since this release makes a breaking change to PRNG key semantics that may require users to update their code, we are bumping the "meso" version of JAX to signify this.

• Breaking changes

  • Enable jax_threefry_partitionable by default (see the update note).

  • This release drops support for Mac x86 wheels. Mac ARM of course remains supported. For a recent discussion, see Will macOS x86 still be supported for the next few years? jax-ml/jax#22936.

    Two key factors motivated this decision:

    • The Mac x86 build (only) has a number of test failures and crashes. We would prefer to ship no release than a broken release.
    • Mac x86 hardware is end-of-life and cannot be easily obtained for developers at this point. So it is difficult for us to fix this kind of problem even if we wanted to.

    We are open to readding support for Mac x86 if the community is willing to help support that platform: in particular, we would need the JAX test suite to pass cleanly on Mac x86 before we could ship releases again.

• Changes:

  • The minimum NumPy version is now 1.25. NumPy 1.25 will remain the minimum supported version until June 2025.
  • The minimum SciPy version is now 1.11. SciPy 1.11 will remain the minimum supported version until June 2025.
  • {func}jax.numpy.einsum now defaults to optimize='auto' rather than optimize='optimal'. This avoids exponentially-scaling trace-time in the case of many arguments ({jax-issue}[#25214](https://github.com/jax-ml/jax/issues/25214)).
  • {func}jax.numpy.linalg.solve no longer supports batched 1D arguments on the right hand side. To recover the previous behavior in these cases, use solve(a, b[..., None]).squeeze(-1).

• New Features

  • {func}jax.numpy.fft.fftn, {func}jax.numpy.fft.rfftn, {func}jax.numpy.fft.ifftn, and {func}jax.numpy.fft.irfftn now support transforms in more than 3 dimensions, which was previously the limit. See {jax-issue}[#25606](https://github.com/jax-ml/jax/issues/25606) for more details.
  • Support added for user defined state in the FFI via the new {func}jax.ffi.register_ffi_type_id function.
  • The AOT lowering .as_text() method now supports the debug_info option to include debugging information, e.g., source location, in the output.

• Deprecations

... (truncated)

Commits

• c25fb92 Release JAX 0.5.0
• a527aba Reverts f1b894d14a28ac22a037fb79177b991275c75a18
• ce85b89 [sharding_in_types] Error out for reshape for splits like this: (4, 6, 8) -...
• 7cac76d Update XLA dependency to use revision
• d3be190 [Mosaic GPU] Delete unused declarations of mosaic_gpu_memcpy_async_h2d.
• d34c40f [mosaic_gpu] Added a serialization pass
• af66719 [sharding_in_types] Rename .at[...].get(out_spec) to `.at[...].get(out_shar...
• 97cd748 Rename out_type -> out_sharding parameter on einsum
• 49224d6 Replace Auto/User/Collective AxisTypes names with Hidden/Visible/Collective.
• bd22bfe [Mosaic TPU] Use large to compact 2nd minor retiling for conversions going bo...
• 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 <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)

Summary by Sourcery

Build:

• Update JAX dependency to version 0.5.0.

[sourcery-ai]

Reviewer's Guide by Sourcery

This pull request updates the jax dependency version requirement from <=0.2.12 to <=0.5.0. The changes are implemented by modifying the version specification in the setup.cfg file across multiple dependency groups to allow using the newer JAX release with added functionalities and breaking changes.

Flow diagram for JAX Dependency Version Update in setup.cfg

flowchart TD
    A[Update JAX dependency in setup.cfg] --> B[Detect dependency groups]
    B --> C[neural_tangents group]
    B --> D[all group]
    C --> E[Old: jax <= 0.2.12]
    C --> F[New: jax <= 0.5.0]
    D --> G[Old: jax <= 0.2.12]
    D --> H[New: jax <= 0.5.0]

    style A fill:#f9f,stroke:#333,stroke-width:2px
    style F fill:#bbf,stroke:#333,stroke-width:2px
    style H fill:#bbf,stroke:#333,stroke-width:2px

Loading

File-Level Changes

Change	Details	Files
Updated the jax dependency version requirement.	Modified jax version in the GBDT dependency group from <=0.2.12 to <=0.5.0 Modified jax version in the all dependency group from <=0.2.12 to <=0.5.0	setup.cfg

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!
• Generate a plan of action for an issue: Comment @sourcery-ai plan on
  an issue to generate a plan of action for it.

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

We have skipped reviewing this pull request. It seems to have been created by a bot (hey, dependabot[bot]!). We assume it knows what it's doing!