pip prod(deps): bump jax from 0.4.38 to 0.5.0 (#3764)

Bumps [jax](https://github.com/jax-ml/jax) from 0.4.38 to 0.5.0.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a
href="https://github.com/jax-ml/jax/releases">jax's
releases</a>.</em></p>
<blockquote>
<h2>JAX v0.5.0</h2>
<p>As of this release, JAX now uses <a
href="https://jax.readthedocs.io/en/latest/jep/25516-effver.html">effort-based
versioning</a>.
Since this release makes a breaking change to PRNG key semantics that
may require users to update their code, we are bumping the
&quot;meso&quot; version of JAX
to signify this.</p>
<ul>
<li>
<p>Breaking changes</p>
<ul>
<li>
<p>Enable <code>jax_threefry_partitionable</code> by default (see
<a href="https://github.com/jax-ml/jax/discussions/18480">the update
note</a>).</p>
</li>
<li>
<p>This release drops support for Mac x86 wheels. Mac ARM of course
remains
supported. For a recent discussion, see <a
href="https://github.com/jax-ml/jax/discussions/22936">https://github.com/jax-ml/jax/discussions/22936</a>.</p>
<p>Two key factors motivated this decision:</p>
<ul>
<li>The Mac x86 build (only) has a number of test failures and crashes.
We
would prefer to ship no release than a broken release.</li>
<li>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.</li>
</ul>
<p>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.</p>
</li>
</ul>
</li>
<li>
<p>Changes:</p>
<ul>
<li>The minimum NumPy version is now 1.25. NumPy 1.25 will remain the
minimum
supported version until June 2025.</li>
<li>The minimum SciPy version is now 1.11. SciPy 1.11 will remain the
minimum
supported version until June 2025.</li>
<li><code>jax.numpy.einsum</code> now defaults to
<code>optimize='auto'</code> rather than
<code>optimize='optimal'</code>. This avoids exponentially-scaling
trace-time in
the case of many arguments
(<code>[#25214](https://github.com/jax-ml/jax/issues/25214)</code>).</li>
<li><code>jax.numpy.linalg.solve</code> no longer supports batched 1D
arguments
on the right hand side. To recover the previous behavior in these cases,
use <code>solve(a, b[..., None]).squeeze(-1)</code>.</li>
</ul>
</li>
<li>
<p>New Features</p>
<ul>
<li><code>jax.numpy.fft.fftn</code>, <code>jax.numpy.fft.rfftn</code>,
<code>jax.numpy.fft.ifftn</code>, and <code>jax.numpy.fft.irfftn</code>
now support
transforms in more than 3 dimensions, which was previously the limit.
See
<code>[#25606](https://github.com/jax-ml/jax/issues/25606)</code> for
more details.</li>
<li>Support added for user defined state in the FFI via the new
<code>jax.ffi.register_ffi_type_id</code> function.</li>
<li>The AOT lowering <code>.as_text()</code> method now supports the
<code>debug_info</code> option
to include debugging information, e.g., source location, in the
output.</li>
</ul>
</li>
<li>
<p>Deprecations</p>
<ul>
<li>From <code>jax.interpreters.xla</code>, <code>abstractify</code> and
<code>pytype_aval_mappings</code>
are now deprecated, having been replaced by symbols of the same name
in <code>jax.core</code>.</li>
</ul>
</li>
</ul>
<!-- raw HTML omitted -->
</blockquote>
<p>... (truncated)</p>
</details>
<details>
<summary>Changelog</summary>
<p><em>Sourced from <a
href="https://github.com/jax-ml/jax/blob/main/CHANGELOG.md">jax's
changelog</a>.</em></p>
<blockquote>
<h2>jax 0.5.0 (Jan 17, 2025)</h2>
<p>As of this release, JAX now uses
<a
href="https://jax.readthedocs.io/en/latest/jep/25516-effver.html">effort-based
versioning</a>.
Since this release makes a breaking change to PRNG key semantics that
may require users to update their code, we are bumping the
&quot;meso&quot; version of JAX
to signify this.</p>
<ul>
<li>
<p>Breaking changes</p>
<ul>
<li>
<p>Enable <code>jax_threefry_partitionable</code> by default (see
<a href="https://github.com/jax-ml/jax/discussions/18480">the update
note</a>).</p>
</li>
<li>
<p>This release drops support for Mac x86 wheels. Mac ARM of course
remains
supported. For a recent discussion, see
<a
href="https://github.com/jax-ml/jax/discussions/22936">https://github.com/jax-ml/jax/discussions/22936</a>.</p>
<p>Two key factors motivated this decision:</p>
<ul>
<li>The Mac x86 build (only) has a number of test failures and crashes.
We
would prefer to ship no release than a broken release.</li>
<li>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.</li>
</ul>
<p>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.</p>
</li>
</ul>
</li>
<li>
<p>Changes:</p>
<ul>
<li>The minimum NumPy version is now 1.25. NumPy 1.25 will remain the
minimum
supported version until June 2025.</li>
<li>The minimum SciPy version is now 1.11. SciPy 1.11 will remain the
minimum
supported version until June 2025.</li>
<li>{func}<code>jax.numpy.einsum</code> now defaults to
<code>optimize='auto'</code> rather than
<code>optimize='optimal'</code>. This avoids exponentially-scaling
trace-time in
the case of many arguments
({jax-issue}<code>[#25214](https://github.com/jax-ml/jax/issues/25214)</code>).</li>
<li>{func}<code>jax.numpy.linalg.solve</code> no longer supports batched
1D arguments
on the right hand side. To recover the previous behavior in these cases,
use <code>solve(a, b[..., None]).squeeze(-1)</code>.</li>
</ul>
</li>
<li>
<p>New Features</p>
<ul>
<li>{func}<code>jax.numpy.fft.fftn</code>,
{func}<code>jax.numpy.fft.rfftn</code>,
{func}<code>jax.numpy.fft.ifftn</code>, and
{func}<code>jax.numpy.fft.irfftn</code> now support
transforms in more than 3 dimensions, which was previously the limit.
See

{jax-issue}<code>[#25606](https://github.com/jax-ml/jax/issues/25606)</code>
for more details.</li>
<li>Support added for user defined state in the FFI via the new
{func}<code>jax.ffi.register_ffi_type_id</code> function.</li>
<li>The AOT lowering <code>.as_text()</code> method now supports the
<code>debug_info</code> option
to include debugging information, e.g., source location, in the
output.</li>
</ul>
</li>
<li>
<p>Deprecations</p>
</li>
</ul>
<!-- raw HTML omitted -->
</blockquote>
<p>... (truncated)</p>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a
href="https://github.com/jax-ml/jax/commit/c25fb92c44d0961f9bef0edaf0cf294d120772ee"><code>c25fb92</code></a>
Release JAX 0.5.0</li>
<li><a
href="https://github.com/jax-ml/jax/commit/a527aba646260c49fd88f421b911296bf5e289ff"><code>a527aba</code></a>
Reverts f1b894d14a28ac22a037fb79177b991275c75a18</li>
<li><a
href="https://github.com/jax-ml/jax/commit/ce85b8988403bdb7c750e60a173dbf5e4cbfc011"><code>ce85b89</code></a>
[sharding_in_types] Error out for reshape for splits like this:
<code>(4, 6, 8)</code> -...</li>
<li><a
href="https://github.com/jax-ml/jax/commit/7cac76d346c3f8862af8af473344f2ba69f7e67b"><code>7cac76d</code></a>
Update XLA dependency to use revision</li>
<li><a
href="https://github.com/jax-ml/jax/commit/d3be190efb58c9c5ec9c0263ad3b20f66b771dd1"><code>d3be190</code></a>
[Mosaic GPU] Delete unused declarations of
<code>mosaic_gpu_memcpy_async_h2d</code>.</li>
<li><a
href="https://github.com/jax-ml/jax/commit/d34c40f6b6371bf1a41e08e047d7138484931421"><code>d34c40f</code></a>
[mosaic_gpu] Added a serialization pass</li>
<li><a
href="https://github.com/jax-ml/jax/commit/af667199dbc1188d1ee0f70cf8922cdb7d64a701"><code>af66719</code></a>
[sharding_in_types] Rename <code>.at[...].get(out_spec)</code> to
`.at[...].get(out_shar...</li>
<li><a
href="https://github.com/jax-ml/jax/commit/97cd748376fe1f26ff0cd8fd9d4d752bed6ce422"><code>97cd748</code></a>
Rename out_type -&gt; out_sharding parameter on einsum</li>
<li><a
href="https://github.com/jax-ml/jax/commit/49224d6cdb4dbae2a8dcc777c295ce6757d5a475"><code>49224d6</code></a>
Replace Auto/User/Collective AxisTypes names with
Hidden/Visible/Collective.</li>
<li><a
href="https://github.com/jax-ml/jax/commit/bd22bfef7135c2c03829111b82ff67605760aee4"><code>bd22bfe</code></a>
[Mosaic TPU] Use large to compact 2nd minor retiling for conversions
going bo...</li>
<li>Additional commits viewable in <a
href="https://github.com/jax-ml/jax/compare/jax-v0.4.38...jax-v0.5.0">compare
view</a></li>
</ul>
</details>
<br />


[![Dependabot compatibility
score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=jax&package-manager=pip&previous-version=0.4.38&new-version=0.5.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores)

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-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

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)


</details>

---------

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: superlopuh <[email protected]>

pyproject.toml

@@ -39,7 +39,7 @@ docs = [
     "mkdocstrings[python]>=0.27.0",
 ]
 gui = ["textual==1.0.0", "pyclip==0.7"]
-jax = ["jax==0.4.38", "numpy==2.2.1"]
+jax = ["jax==0.5.0", "numpy==2.2.1"]
 riscv = ["riscemu==2.2.7"]
 
 [project.urls]