Bump the jax group across 1 directory with 2 updates (#1501)

Updates the requirements on [jax](https://github.com/jax-ml/jax) and
[diffrax](https://github.com/patrick-kidger/diffrax) to permit the
latest version.
Updates `jax` 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>
<h2>jax 0.4.37 (Dec 9, 2024)</h2>
<p>This is a patch release of jax 0.4.36. Only &quot;jax&quot; was
released at this version.</p>
<ul>
<li>Bug fixes
<ul>
<li>Fixed a bug where <code>jit</code> would error if an argument was
named <code>f</code> (<a
href="https://redirect.github.com/jax-ml/jax/issues/25329">#25329</a>).</li>
<li>Fix a bug that will throw <code>index out of range</code> error in
{func}<code>jax.lax.while_loop</code> if the user register pytree node
class with
different aux data for the flatten and flatten_with_path.</li>
<li>Pinned a new libtpu release (0.0.6) that fixes a compiler bug on TPU
v6e.</li>
</ul>
</li>
</ul>
<h2>jax 0.4.36 (Dec 5, 2024)</h2>
<ul>
<li>Breaking Changes
<ul>
<li>
<p>This release lands &quot;stackless&quot;, an internal change to JAX's
tracing
machinery. We made trace dispatch purely a function of context rather
than a
function of both context and data. This let us delete a lot of machinery
for
managing data-dependent tracing: levels, sublevels,
<code>post_process_call</code>,
<code>new_base_main</code>, <code>custom_bind</code>, and so on. The
change should only affect
users that use JAX internals.</p>
<p>If you do use JAX internals then you may need to
update your code (see
<a
href="https://github.com/jax-ml/jax/commit/c36e1f7c1ad4782060cbc8e8c596d85dfb83986f">https://github.com/jax-ml/jax/commit/c36e1f7c1ad4782060cbc8e8c596d85dfb83986f</a></p>
</li>
</ul>
</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/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.24...jax-v0.4.38">compare
view</a></li>
</ul>
</details>
<br />

Updates `diffrax` to 0.6.2
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a
href="https://github.com/patrick-kidger/diffrax/releases">diffrax's
releases</a>.</em></p>
<blockquote>
<h2>Diffrax v0.6.2</h2>
<ul>
<li>Updated to silence warnings from JAX 0.4.38 (<code>jax.core.Jaxpr
-&gt; jax.extend.core.Jaxpr</code>) (Thanks <a
href="https://github.com/jakevdp"><code>@​jakevdp</code></a>! <a
href="https://redirect.github.com/patrick-kidger/diffrax/issues/544">#544</a>,
<a
href="https://redirect.github.com/patrick-kidger/diffrax/issues/546">#546</a>)</li>
<li>Added <code>diffrax.ForwardMode</code> for efficient forward-mode
autodifferentiation through <code>diffeqsolve</code>. Updated docs to
reflect that there are now two main recommended choices, namely
<code>ForwardMode</code> for forward-mode and
<code>RecursiveCheckpointAdjoint</code> for reverse-mode. (Thanks <a
href="https://github.com/johannahaffner"><code>@​johannahaffner</code></a>!
<a
href="https://redirect.github.com/patrick-kidger/diffrax/issues/537">#537</a>)</li>
</ul>
<p><strong>Full Changelog</strong>: <a
href="https://github.com/patrick-kidger/diffrax/compare/v0.6.1...v0.6.2">https://github.com/patrick-kidger/diffrax/compare/v0.6.1...v0.6.2</a></p>
</blockquote>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/7474749b6a37d4cb74c493b07924f7a3cf361f1f"><code>7474749</code></a>
Bump test versions</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/25898761f9e6e4bfbfc14c561d14d56c78377ad4"><code>2589876</code></a>
Silence warning from pure_callback(vectorized=True)</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/10bff5abfd8e48b859080a9b6408e797dc1d5c7e"><code>10bff5a</code></a>
Version bump</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/da2865e9fac588b59f02cc11c9945613f16499c2"><code>da2865e</code></a>
Update documentation on adjoint methods to better emphasise what should
and s...</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/d28358a7a3e8b1255fed1849be0bb746615b232b"><code>d28358a</code></a>
Forward mode &quot;adjoint&quot; (<a
href="https://redirect.github.com/patrick-kidger/diffrax/issues/537">#537</a>)</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/3553ae1360acbaf9f01143cab300dd0864e85a9e"><code>3553ae1</code></a>
Fix pyright failures from numpy 2.</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/613dd0ea9c6e637e780a7e3ea52920b88fc85d0e"><code>613dd0e</code></a>
Import Jaxpr and ClosedJaxpr from jax.exend</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/ac9f3ec7ee232c34f8c67c6f0073fc70302a457f"><code>ac9f3ec</code></a>
Fix CI</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/d7fc34845bb5782988f1c3308292723c70f25338"><code>d7fc348</code></a>
Fix docs?</li>
<li><a
href="https://github.com/patrick-kidger/diffrax/commit/78531fa2ae15f0e8ce7356148642122a16d3531a"><code>78531fa</code></a>
Bump minimum Equinox version to one that is compatible with latest
JAX</li>
<li>Additional commits viewable in <a
href="https://github.com/patrick-kidger/diffrax/compare/v0.4.1...v0.6.2">compare
view</a></li>
</ul>
</details>
<br />


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 <dependency name> major version` will close this
group update PR and stop Dependabot creating any more for the specific
dependency's major version (unless you unignore this specific
dependency's major version or upgrade to it yourself)
- `@dependabot ignore <dependency name> minor version` will close this
group update PR and stop Dependabot creating any more for the specific
dependency's minor version (unless you unignore this specific
dependency's minor version or upgrade to it yourself)
- `@dependabot ignore <dependency name>` will close this group update PR
and stop Dependabot creating any more for the specific dependency
(unless you unignore this specific dependency or upgrade to it yourself)
- `@dependabot unignore <dependency name>` will remove all of the ignore
conditions of the specified dependency
- `@dependabot unignore <dependency name> <ignore condition>` will
remove the ignore condition of the specified dependency and ignore
conditions


</details>

requirements.txt

@@ -1,6 +1,6 @@
-jax >= 0.4.24, != 0.4.36, <= 0.4.37
+jax >= 0.4.24, != 0.4.36, <= 0.4.38
 colorama <= 0.4.6
-diffrax >= 0.4.1, <= 0.6.1
+diffrax >= 0.4.1, <= 0.6.2
 h5py >= 3.0.0, <= 3.12.1
 interpax >= 0.3.3, <= 0.3.4
 matplotlib >= 3.5.0, <= 3.9.3