Update jax[cpu] requirement from <=0.4.36,>=0.3.2 to >=0.3.2,<=0.4.37 (…

…#23)

Updates the requirements on [jax[cpu]](https://github.com/jax-ml/jax) to
permit the latest version.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a
href="https://github.com/jax-ml/jax/releases">jax[cpu]'s
releases</a>.</em></p>
<blockquote>
<h2>JAX v0.4.37</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
<code>jax.lax.while_loop</code> if the user registers 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>
</blockquote>
</details>
<details>
<summary>Changelog</summary>
<p><em>Sourced from <a
href="https://github.com/jax-ml/jax/blob/main/CHANGELOG.md">jax[cpu]'s
changelog</a>.</em></p>
<blockquote>
<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>
for clues about how to do this). There might also be version skew
issues with JAX libraries that do this. If you find this change breaks
your
non-JAX-internals-using code then try the
<code>config.jax_data_dependent_tracing_fallback</code> flag as a
workaround, and if
you need help updating your code then please file a bug.</p>
</li>
<li>
<p>{func}<code>jax.experimental.jax2tf.convert</code> with
<code>native_serialization=False</code>
or with <code>enable_xla=False</code> have been deprecated since July
2024, with
JAX version 0.4.31. Now we removed support for these use cases.
<code>jax2tf</code>
with native serialization will still be supported.</p>
</li>
<li>
<p>In <code>jax.interpreters.xla</code>, the <code>xb</code>,
<code>xc</code>, and <code>xe</code> symbols have been removed
after being deprecated in JAX v0.4.31. Instead use <code>xb =
jax.lib.xla_bridge</code>,
<code>xc = jax.lib.xla_client</code>, and <code>xe =
jax.lib.xla_extension</code>.</p>
</li>
<li>
<p>The deprecated module <code>jax.experimental.export</code> has been
removed. It was replaced
by {mod}<code>jax.export</code> in JAX v0.4.30. See the <a
href="https://jax.readthedocs.io/en/latest/export/export.html#migration-guide-from-jax-experimental-export">migration
guide</a>
for information on migrating to the new API.</p>
</li>
<li>
<p>The <code>initial</code> argument to
{func}<code>jax.nn.softmax</code> and
{func}<code>jax.nn.log_softmax</code>
has been removed, after being deprecated in v0.4.27.</p>
</li>
<li>
<p>Calling <code>np.asarray</code> on typed PRNG keys (i.e. keys
produced by :func:<code>jax.random.key</code>)
now raises an error. Previously, this returned a scalar object
array.</p>
</li>
<li>
<p>The following deprecated methods and functions in
{mod}<code>jax.export</code> have
been removed:</p>
<ul>
<li><code>jax.export.DisabledSafetyCheck.shape_assertions</code>: it had
no effect
already.</li>
<li><code>jax.export.Exported.lowering_platforms</code>: use
<code>platforms</code>.</li>
<li><code>jax.export.Exported.mlir_module_serialization_version</code>:
use <code>calling_convention_version</code>.</li>
</ul>
</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/ffb07cdadb5dc3bc43485cf041dbc2b43136109e"><code>ffb07cd</code></a>
Update versions for v0.4.37 release.</li>
<li><a
href="https://github.com/jax-ml/jax/commit/95892fdac86524151b6dadd7d8bedbf915f1500f"><code>95892fd</code></a>
Use private names for args in api_util to avoid shadowing kwargs
keys.</li>
<li><a
href="https://github.com/jax-ml/jax/commit/65b60884114261549ffc2eb937162bdeaa493928"><code>65b6088</code></a>
Avoid index out of range error in carry structure check</li>
<li><a
href="https://github.com/jax-ml/jax/commit/259194a69f52a06847a9ff11eb268072e91fd65f"><code>259194a</code></a>
[Pallas] Fix shard_axis in dma_start interpret mode rule.</li>
<li><a
href="https://github.com/jax-ml/jax/commit/7e6620a57775084dfa8d438ae4fd27f3ef365018"><code>7e6620a</code></a>
JAX release 0.4.36.</li>
<li><a
href="https://github.com/jax-ml/jax/commit/23d5c10ff0704f66ad7ec65a8cdcd09bd2420591"><code>23d5c10</code></a>
[Mosaic:TPU] Fix fully replicated relayout</li>
<li><a
href="https://github.com/jax-ml/jax/commit/2a4a0e8d6fb36b59f9c6f24e0018d42c8c8d8ee9"><code>2a4a0e8</code></a>
[jax:custom_partitioning] Implement SdyShardingRule to support</li>
<li><a
href="https://github.com/jax-ml/jax/commit/f73fa7a7ad64b2f15e8669beed14600704287b93"><code>f73fa7a</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/25290">#25290</a>
from jakevdp:reduction-where</li>
<li><a
href="https://github.com/jax-ml/jax/commit/a71f9a62e6f67640a4b0578d042b07792fcf407a"><code>a71f9a6</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/25271">#25271</a>
from jakevdp:fix-vector-norm</li>
<li><a
href="https://github.com/jax-ml/jax/commit/e20a483befbb80bbf782b931ec57a44c78c313b8"><code>e20a483</code></a>
[JAX] Add end-to-end execution support in colocated Python API</li>
<li>Additional commits viewable in <a
href="https://github.com/jax-ml/jax/compare/jax-v0.3.2...jax-v0.4.37">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 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>

requirements-dev.txt

@@ -1,7 +1,6 @@
 # standard install requirements
 -r ./requirements.txt
 
-jax[cpu] >= 0.3.2, <= 0.4.36
 scipy >= 1.10.0, <= 1.14.1
 
 # building the docs