pip prod(deps): bump jax from 0.4.37 to 0.4.38 (#3651)

Bumps [jax](https://github.com/jax-ml/jax) from 0.4.37 to 0.4.38.
<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.4.38</h2>
<ul>
<li>
<p>Changes:</p>
<ul>
<li><code>jax.tree.flatten_with_path</code> and
<code>jax.tree.map_with_path</code> are added
as shortcuts of the corresponding <code>tree_util</code> functions.</li>
</ul>
</li>
<li>
<p>Deprecations</p>
<ul>
<li>a number of APIs in the internal <code>jax.core</code> namespace
have been deprecated.
Most were no-ops, were little-used, or can be replaced by APIs of the
same
name in <code>jax.extend.core</code>; see the documentation for
{mod}<code>jax.extend</code>
for information on the compatibility guarantees of these semi-public
extensions.</li>
<li>Several previously-deprecated APIs have been removed, including:
<ul>
<li>from <code>jax.core</code>: <code>check_eqn</code>,
<code>check_type</code>, <code>check_valid_jaxtype</code>, and
<code>non_negative_dim</code>.</li>
<li>from <code>jax.lib.xla_bridge</code>: <code>xla_client</code> and
<code>default_backend</code>.</li>
<li>from <code>jax.lib.xla_client</code>: <code>_xla</code> and
<code>bfloat16</code>.</li>
<li>from <code>jax.numpy</code>: <code>round_</code>.</li>
</ul>
</li>
</ul>
</li>
<li>
<p>New Features</p>
<ul>
<li><code>jax.export.export</code> can be used for device-polymorphic
export with
shardings constructed with {func}<code>jax.sharding.AbstractMesh</code>.
See the <a
href="https://jax.readthedocs.io/en/latest/export/export.html#device-polymorphic-export">jax.export
documentation</a>.</li>
<li>Added <code>jax.lax.split</code>. This is a primitive version of
<code>jax.numpy.split</code>, added because it yields a more compact
transpose during automatic differentiation.</li>
</ul>
</li>
</ul>
</blockquote>
</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.4.38 (Dec 17, 2024)</h2>
<ul>
<li>
<p>Changes:</p>
<ul>
<li><code>jax.tree.flatten_with_path</code> and
<code>jax.tree.map_with_path</code> are added
as shortcuts of the corresponding <code>tree_util</code> functions.</li>
</ul>
</li>
<li>
<p>Deprecations</p>
<ul>
<li>a number of APIs in the internal <code>jax.core</code> namespace
have been deprecated.
Most were no-ops, were little-used, or can be replaced by APIs of the
same
name in {mod}<code>jax.extend.core</code>; see the documentation for
{mod}<code>jax.extend</code>
for information on the compatibility guarantees of these semi-public
extensions.</li>
<li>Several previously-deprecated APIs have been removed, including:
<ul>
<li>from {mod}<code>jax.core</code>: <code>check_eqn</code>,
<code>check_type</code>, <code>check_valid_jaxtype</code>, and
<code>non_negative_dim</code>.</li>
<li>from {mod}<code>jax.lib.xla_bridge</code>: <code>xla_client</code>
and <code>default_backend</code>.</li>
<li>from {mod}<code>jax.lib.xla_client</code>: <code>_xla</code> and
<code>bfloat16</code>.</li>
<li>from {mod}<code>jax.numpy</code>: <code>round_</code>.</li>
</ul>
</li>
</ul>
</li>
<li>
<p>New Features</p>
<ul>
<li>{func}<code>jax.export.export</code> can be used for
device-polymorphic export with
shardings constructed with {func}<code>jax.sharding.AbstractMesh</code>.
See the <a
href="https://jax.readthedocs.io/en/latest/export/export.html#device-polymorphic-export">jax.export
documentation</a>.</li>
<li>Added {func}<code>jax.lax.split</code>. This is a primitive version
of
{func}<code>jax.numpy.split</code>, added because it yields a more
compact
transpose during automatic differentiation.</li>
</ul>
</li>
</ul>
</blockquote>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a
href="https://github.com/jax-ml/jax/commit/c4e3f5a3563e4ef34ba8c4dfa797ca4a47c17188"><code>c4e3f5a</code></a>
Release v0.4.38</li>
<li><a
href="https://github.com/jax-ml/jax/commit/0fa541972eef5609c1234543cf99301d1601231c"><code>0fa5419</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/25456">#25456</a>
from jakevdp:xla-abstractify</li>
<li><a
href="https://github.com/jax-ml/jax/commit/7fe25791e4d79ea21255a6ec177484dc7d485656"><code>7fe2579</code></a>
Update XLA dependency to use revision</li>
<li><a
href="https://github.com/jax-ml/jax/commit/7bbba40cd42aa47a6880308469b56294c1f3512c"><code>7bbba40</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/25473">#25473</a>
from adamjstewart:build/no-git</li>
<li><a
href="https://github.com/jax-ml/jax/commit/7de9eb20dfb57c6e4a0d7d9abadb20c70af4600d"><code>7de9eb2</code></a>
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0</li>
<li><a
href="https://github.com/jax-ml/jax/commit/4911a396b2fdf9ffaf5b6c13f31059f7954f1ae1"><code>4911a39</code></a>
[Mosaic TPU] Add support for the interleaved pack format to
tpu.unpack_subele...</li>
<li><a
href="https://github.com/jax-ml/jax/commit/36b12d58f4df21d5cce629bda7d5c48280c35360"><code>36b12d5</code></a>
[Mosaic GPU] Add end-to-end lowering example for a pointwise kernel
using the...</li>
<li><a
href="https://github.com/jax-ml/jax/commit/2c722d9b13c65eaee5bf9a55f5186b10ea772477"><code>2c722d9</code></a>
Cleanup: toward merging core.concrete_aval &amp; xla.abstractify</li>
<li><a
href="https://github.com/jax-ml/jax/commit/473e2bf52710e89fcd3ad2288afe877b140e926c"><code>473e2bf</code></a>
Put abstract_mesh on every eqn so that we can preserve it during
<code>eval_jaxpr</code>...</li>
<li><a
href="https://github.com/jax-ml/jax/commit/e1f037bffd6420d64af47aa853e7d2347c9e8004"><code>e1f037b</code></a>
Reduce the max_examples for splash attention tests</li>
<li>Additional commits viewable in <a
href="https://github.com/jax-ml/jax/compare/jax-v0.4.37...jax-v0.4.38">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.37&new-version=0.4.38)](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: Sasha Lopoukhine <[email protected]>

pyproject.toml

@@ -34,7 +34,7 @@ dev = [
     "pyright==1.1.391",
 ]
 gui = ["textual==1.0.0", "pyclip==0.7"]
-jax = ["jax==0.4.37", "numpy==2.2.1"]
+jax = ["jax==0.4.38", "numpy==2.2.1"]
 onnx = ["onnx==1.17.0", "numpy==2.2.1"]
 riscv = ["riscemu==2.2.7"]
 wgpu = ["wgpu==0.19.3"]