From acf157ac7f7f56aba3477b1cda7b6adb67617407 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 20 Jan 2025 09:58:47 +0100
Subject: [PATCH] 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] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: superlopuh <superlopuh@users.noreply.github.com>
---
 pyproject.toml |  2 +-
 uv.lock        | 46 +++++++++++++++++++++-------------------------
 2 files changed, 22 insertions(+), 26 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 04d8052298..8fc316fad5 100644
--- a/pyproject.toml
+++ b/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]
diff --git a/uv.lock b/uv.lock
index dbe9c71f74..295155d999 100644
--- a/uv.lock
+++ b/uv.lock
@@ -835,7 +835,7 @@ wheels = [
 
 [[package]]
 name = "jax"
-version = "0.4.38"
+version = "0.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jaxlib" },
@@ -844,14 +844,14 @@ dependencies = [
     { name = "opt-einsum" },
     { name = "scipy" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/fb/e5/c4aa9644bb96b7f6747bd7c9f8cda7665ca5e194fa2542b2dea3ff730701/jax-0.4.38.tar.gz", hash = "sha256:43bae65881628319e0a2148e8f81a202fbc2b8d048e35c7cb1df2416672fa4a8", size = 1930034 }
+sdist = { url = "https://files.pythonhosted.org/packages/4a/cb/22d62b26284f08e62d6eb64603d3b010004cfdb7a97ce6cca5c6cf86edab/jax-0.5.0.tar.gz", hash = "sha256:49df70bf293a345a7fb519f71193506d37a024c4f850b358042eb32d502c81c8", size = 1959707 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/22/49/b4418a7a892c0dd64442bbbeef54e1cdfe722dfc5a7bf0d611d3f5f90e99/jax-0.4.38-py3-none-any.whl", hash = "sha256:78987306f7041ea8500d99df1a17c33ed92620c2268c4c3677fb24e06712be64", size = 2236864 },
+    { url = "https://files.pythonhosted.org/packages/f4/58/cc0721a1030fcbab0984beea0bf3c4610ec103f738423cdfa9c4ceb40598/jax-0.5.0-py3-none-any.whl", hash = "sha256:b3907aa87ae2c340b39cdbf80c07a74550369cafcaf7398fb60ba58d167345ab", size = 2270365 },
 ]
 
 [[package]]
 name = "jaxlib"
-version = "0.4.38"
+version = "0.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "ml-dtypes" },
@@ -859,26 +859,22 @@ dependencies = [
     { name = "scipy" },
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/ee/d4/e6a0881a88b8f17491c2ee271fd77c348b0221d9e2ec92dad23a2c9e41bc/jaxlib-0.4.38-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:55c19b9d3f33a6fc59f644aa5a21fba02639ccdd776cb4a9b5526625f57839ff", size = 99663603 },
-    { url = "https://files.pythonhosted.org/packages/b6/6d/11569ce873f04c82ec22e58d822f4187dccae1d400c0d6dd05ed314d5328/jaxlib-0.4.38-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:30b2f52cb50d74734af2f477c2533a7a583e3bb7b2c8acdeb361ee77d940577a", size = 79475708 },
-    { url = "https://files.pythonhosted.org/packages/72/61/1de2405d13089c83b1ad87ec0266479c9d00080659dae2474892ae356306/jaxlib-0.4.38-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:ee19c163a8fdf0839d4c18b88a5fbfb4e731ba7c437416d3e5483e570bb764e4", size = 93219045 },
-    { url = "https://files.pythonhosted.org/packages/9c/24/0829decf233c6af9efe7c53888ae8ac72395e0979869cd9cee487e35dac3/jaxlib-0.4.38-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:61aeccb9a27c67fdb8450f6357240019cd4511cb9d62a44e4764756d384853ad", size = 101732107 },
-    { url = "https://files.pythonhosted.org/packages/0d/04/120c4caac6151f7297fedf9dd776362aa2d417d3f87bda826050b4da45e8/jaxlib-0.4.38-cp310-cp310-win_amd64.whl", hash = "sha256:d6ab745a89d0fb737a36fe1d8b86659e3fffe6ee8303b20651b26193d5edc0ef", size = 64223924 },
-    { url = "https://files.pythonhosted.org/packages/b0/6a/b9fba73eb5e758e40a514919e096a039d27dc0ab4776a6cc977f5153a55f/jaxlib-0.4.38-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:b67fdeabd6dfed08b7768f3bdffb521160085f8305669bd197beef61d08de08b", size = 99679916 },
-    { url = "https://files.pythonhosted.org/packages/44/2a/3458130d44d44038fd6974e7c43948f68408f685063203b82229b9b72c1a/jaxlib-0.4.38-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3fb0eaae7369157afecbead50aaf29e73ffddfa77a2335d721bd9794f3c510e4", size = 79488377 },
-    { url = "https://files.pythonhosted.org/packages/94/96/7d9a0b9f35af4727df44b68ade4c6f15163840727d1cb47251b1ea515e30/jaxlib-0.4.38-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:43db58c4c427627296366a56c10318e1f00f503690e17f94bb4344293e1995e0", size = 93241543 },
-    { url = "https://files.pythonhosted.org/packages/a3/2d/68f85037e60c981b37b18b23ace458c677199dea4722ddce541b48ddfc63/jaxlib-0.4.38-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:2751ff7037d6a997d0be0e77cc4be381c5a9f9bb8b314edb755c13a6fd969f45", size = 101751923 },
-    { url = "https://files.pythonhosted.org/packages/cc/24/a9c571c8a189f58e0b54b14d53fc7f5a0a06e4f1d7ab9edcf8d1d91d07e7/jaxlib-0.4.38-cp311-cp311-win_amd64.whl", hash = "sha256:35226968fc9de6873d1571670eac4117f5ed80e955f7a1775204d1044abe16c6", size = 64255189 },
-    { url = "https://files.pythonhosted.org/packages/49/df/08b94c593c0867c7eaa334592807ba74495de4be90580f360db8b96221dc/jaxlib-0.4.38-cp312-cp312-macosx_10_14_x86_64.whl", hash = "sha256:3fefea985f0415816f3bbafd3f03a437050275ef9bac9a72c1314e1644ac57c1", size = 99737849 },
-    { url = "https://files.pythonhosted.org/packages/ab/b1/c9d2a7ba9ebeabb7ac37082f4c466364f475dc7550a79358c0f0aa89fdf2/jaxlib-0.4.38-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:f33bcafe32c97a562ecf6894d7c41674c80c0acdedfa5423d49af51147149874", size = 79509242 },
-    { url = "https://files.pythonhosted.org/packages/53/25/dd670d8bdf3799ece76d12cfe6a6a250ea256057aa4b0fcace4753a99d2d/jaxlib-0.4.38-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:496f45b0e001a2341309cd0c74af0b670537dced79c168cb230cfcc773f0aa86", size = 93251503 },
-    { url = "https://files.pythonhosted.org/packages/f9/cc/37fce5162f6b9070203fd76cc0f298d9b3bfdf01939a78935a6078d63621/jaxlib-0.4.38-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:dad6c0a96567c06d083c0469fec40f201210b099365bd698be31a6d2ec88fd59", size = 101792792 },
-    { url = "https://files.pythonhosted.org/packages/6f/7a/8515950a60a4ea5b13cc98fc0a42e36553b2db5a6eedc00d3bd7836f77b5/jaxlib-0.4.38-cp312-cp312-win_amd64.whl", hash = "sha256:966cdec36cfa978f5b4582bcb4147fe511725b94c1a752dac3a5f52ce46b6fa3", size = 64288223 },
-    { url = "https://files.pythonhosted.org/packages/91/03/aee503c7077c6dbbd568842303426c6ec1cef9bff330c418c9e71906cccd/jaxlib-0.4.38-cp313-cp313-macosx_10_14_x86_64.whl", hash = "sha256:41e55ae5818a882e5789e848f6f16687ac132bcfbb5a5fa114a5d18b78d05f2d", size = 99739026 },
-    { url = "https://files.pythonhosted.org/packages/cb/bf/fbbf61da319611d88e11c691d5a2077039208ded05e1731dea940f824a59/jaxlib-0.4.38-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:6fe326b8af366387dd47ccf312583b2b17fed12712c9b74a648b18a13cbdbabf", size = 79508735 },
-    { url = "https://files.pythonhosted.org/packages/e4/0b/8cbff0b6d62a4694351c49baf53b7ed8deb8a6854d129408c38158e11676/jaxlib-0.4.38-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:248cca3771ebf24b070f49701364ceada33e6139445b06c782cca5ac5ad92bf4", size = 93251882 },
-    { url = "https://files.pythonhosted.org/packages/15/57/7f0283273b69c417071bcd2f4c2ed076479ec5ffc22a647f13c21da8d071/jaxlib-0.4.38-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:2ce77ba8cda9259a4bca97afc1c722e4291a6c463a63f8d372c6edc85117d625", size = 101791137 },
-    { url = "https://files.pythonhosted.org/packages/de/de/d6c4d234cd426b97459cb070af90792b48643967a0d28641379ee9e10fc9/jaxlib-0.4.38-cp313-cp313-win_amd64.whl", hash = "sha256:4103db0b3a38a5dc132741237453c24d8547290a22079ba1b577d6c88c95300a", size = 64288459 },
+    { url = "https://files.pythonhosted.org/packages/c8/41/3e4ac64df72c4da126df3fd66a2214025a46b6263f7be266728e7b8e473e/jaxlib-0.5.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1b8a6c4345f137f387650de2dbc488c20251b7412b55dd648e1a4f13bcf507fb", size = 79248968 },
+    { url = "https://files.pythonhosted.org/packages/1e/5f/2a16e61f1d54ae5f55fbf3cb3e22ef5bb01bf9d7d6474e0d34fedba19c4d/jaxlib-0.5.0-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:5b2efe3dfebf18a84c451d3803ac884ee242021c1113b279c13f4bbc378c3dc0", size = 93181077 },
+    { url = "https://files.pythonhosted.org/packages/08/c3/573e2f01b99f1247e8fbe1aa46b95a0faa68ef208f9a8e8ef775d607b3e6/jaxlib-0.5.0-cp310-cp310-manylinux2014_x86_64.whl", hash = "sha256:74440b632107336400d4f97a16481d767f13ea914c53ba14e544c6fda54819b3", size = 101969119 },
+    { url = "https://files.pythonhosted.org/packages/6e/38/512f61ea13da41ca47f2411d7c05af0cf74a37f225e16725ed0e6fb58893/jaxlib-0.5.0-cp310-cp310-win_amd64.whl", hash = "sha256:53478a28eee6c2ef01759b05a9491702daef9268c3ed013d6f8e2e5f5cae0887", size = 63883394 },
+    { url = "https://files.pythonhosted.org/packages/92/4b/8875870ff52ad3fbea876c905228f691f05c8dc8556b226cbfaf0fba7f62/jaxlib-0.5.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6cd762ed1623132499fa701c4203446102e0a9c82ca23194b87288f746d12a29", size = 79242870 },
+    { url = "https://files.pythonhosted.org/packages/a0/0f/00cdfa411d7218e4696c10c5867f7d3c396219adbcaeb02e95108ca802de/jaxlib-0.5.0-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:63088dbfaa85bb56cd521a925a3472fd7328b18ec93c2d8ffa85af331095c995", size = 93181807 },
+    { url = "https://files.pythonhosted.org/packages/58/8e/a5c29db03d5a93b0326e297b556d0e0a9805e9c9c1ae5f82f69557273faa/jaxlib-0.5.0-cp311-cp311-manylinux2014_x86_64.whl", hash = "sha256:09113ef1582ba34d7cbc440fedb318f4855b59b776711a8aba2473c9727d3025", size = 101969212 },
+    { url = "https://files.pythonhosted.org/packages/70/86/ceae20e4f37fa07f1cc95551cc0f49170d0db46d2e82fdf511d26bffd801/jaxlib-0.5.0-cp311-cp311-win_amd64.whl", hash = "sha256:78289fc3ddc1e4e9510de2536a6375df9fe1c50de0ac60826c286b7a5c5090fe", size = 63881994 },
+    { url = "https://files.pythonhosted.org/packages/57/d6/d971b40cb156e0637aa3c1522a1e803b641142e9a8f3ade6a574711bb073/jaxlib-0.5.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:73e335715760c56e635109d61426435a5d7f46f3363a115daea09427d5cd0efd", size = 79246087 },
+    { url = "https://files.pythonhosted.org/packages/41/2e/ba9770330077c3e4082cd0353e6a61419f79bff3e2197f904ce70167b9ad/jaxlib-0.5.0-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:4b4b01afb0ddec96c08356bff2bb685ddbe97fdffe4ed6e2d834b30aba972f22", size = 93179593 },
+    { url = "https://files.pythonhosted.org/packages/66/e9/211ba3e46ec22c722c4d61a739cfccf79b0618006d6f5fa53eb4eb93ed6d/jaxlib-0.5.0-cp312-cp312-manylinux2014_x86_64.whl", hash = "sha256:f980c733e98c998a8da87c9a8cc61b6726d0be667a58bd664c1d717b4b4eae75", size = 101984785 },
+    { url = "https://files.pythonhosted.org/packages/2d/cb/11bb92324afb6ba678f388e10b78d6b02196bc8887eb5aa0d85ce398edf9/jaxlib-0.5.0-cp312-cp312-win_amd64.whl", hash = "sha256:5baedbeeb60fa493c7528783254f04c6e986a2826266b198ed37e9336af2ef8c", size = 63899871 },
+    { url = "https://files.pythonhosted.org/packages/22/ac/e400473e6a2f405fd6e4dc40a713bb9a3868a3f76a8ffc5eb66f6e686002/jaxlib-0.5.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ed18ea7161d03aa8fd4d1b55494882f21420efdfea68e5f298c4aebcf2ac3f34", size = 79245359 },
+    { url = "https://files.pythonhosted.org/packages/44/2d/c210abf4a9b2ce2e0858fcd3567c8773a739114e37d751af6c228901af57/jaxlib-0.5.0-cp313-cp313-manylinux2014_aarch64.whl", hash = "sha256:7d9b17a7ea19355d45ecdb2ff0db5d707a86f0c5a862d94b89b4568d6c45311a", size = 93180025 },
+    { url = "https://files.pythonhosted.org/packages/30/f8/316f7b4797c5eb50c6d70e461724a7cbe08b4505ca4da1bfd260c135895a/jaxlib-0.5.0-cp313-cp313-manylinux2014_x86_64.whl", hash = "sha256:11eef01d37c0f1c5306265b76f207f1002d13480ded2e31fd63ec76912c93ca2", size = 101982281 },
+    { url = "https://files.pythonhosted.org/packages/4c/f2/cfa012a0417c9b13b44c8e1d3ebf5fd04e8bb738b5c93e20c9fc97919880/jaxlib-0.5.0-cp313-cp313-win_amd64.whl", hash = "sha256:61b4d26cd6a0c49ba0b1e4340c7d29198913ee2dc70b65ee90752717d22305bb", size = 63900219 },
 ]
 
 [[package]]
@@ -2899,7 +2895,7 @@ requires-dist = [
     { name = "filecheck", marker = "extra == 'dev'", specifier = "==1.0.1" },
     { name = "immutabledict", specifier = "<4.2.2" },
     { name = "ipykernel", marker = "extra == 'dev'" },
-    { name = "jax", marker = "extra == 'jax'", specifier = "==0.4.38" },
+    { name = "jax", marker = "extra == 'jax'", specifier = "==0.5.0" },
     { name = "lit", marker = "extra == 'dev'", specifier = "<19.0.0" },
     { name = "marimo", marker = "extra == 'dev'", specifier = "==0.10.14" },
     { name = "mkdocs", marker = "extra == 'docs'", specifier = ">=1.6.1" },