pip prod(deps): bump jax from 0.4.31 to 0.4.32 (#3161)

Bumps [jax](https://github.com/google/jax) from 0.4.31 to 0.4.32.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a
href="https://github.com/google/jax/releases">jax's
releases</a>.</em></p>
<blockquote>
<h2>JAX release v0.4.32</h2>
<ul>
<li>
<p>New Functionality</p>
<ul>
<li>Added <code>jax.extend.ffi.ffi_call</code> and
<code>jax.extend.ffi.ffi_lowering</code>
to support the use of the new <code>ffi-tutorial</code> to interface
with custom
C++ and CUDA code from JAX.</li>
</ul>
</li>
<li>
<p>Changes</p>
<ul>
<li><code>jax_enable_memories</code> flag is set to <code>True</code> by
default.</li>
<li><code>jax.numpy</code> now supports v2023.12 of the Python Array API
Standard.
See <code>python-array-api</code> for more information.</li>
<li>Computations on the CPU backend may now be dispatched asynchronously
in
more cases. Previously non-parallel computations were always dispatched
synchronously. You can recover the old behavior by setting
<code>jax.config.update('jax_cpu_enable_async_dispatch',
False)</code>.</li>
<li>Added new <code>jax.process_indices</code> function to replace the
<code>jax.host_ids()</code> function that was deprecated in JAX
v0.2.13.</li>
<li>To align with the behavior of <code>numpy.fabs</code>,
<code>jax.numpy.fabs</code> has been
modified to no longer support <code>complex dtypes</code>.</li>
<li><code>jax.tree_util.register_dataclass</code> now checks that
<code>data_fields</code>
and <code>meta_fields</code> includes all dataclass fields with
<code>init=True</code>
and only them, if <code>nodetype</code> is a dataclass.</li>
<li>Several <code>jax.numpy</code> functions now have full
<code>jax.numpy.ufunc</code>
interfaces, including <code>jax.numpy.add</code>,
<code>jax.numpy.multiply</code>,
<code>jax.numpy.bitwise_and</code>, <code>jax.numpy.bitwise_or</code>,
<code>jax.numpy.bitwise_xor</code>, <code>jax.numpy.logical_and</code>,
<code>jax.numpy.logical_and</code>, and
<code>jax.numpy.logical_and</code>.
In future releases we plan to expand these to other ufuncs.</li>
<li>Added <code>jax.lax.optimization_barrier</code>, which allows users
to prevent
compiler optimizations such as common-subexpression elimination and to
control scheduling.</li>
</ul>
</li>
<li>
<p>Breaking changes</p>
<ul>
<li>The MHLO MLIR dialect (<code>jax.extend.mlir.mhlo</code>) has been
removed. Use the
<code>stablehlo</code> dialect instead.</li>
</ul>
</li>
<li>
<p>Deprecations</p>
<ul>
<li>Complex inputs to <code>jax.numpy.clip</code> and
{func}<code>jax.numpy.hypot</code> are
no longer allowed, after being deprecated since JAX v0.4.27.</li>
<li>Deprecated the following APIs:
<ul>
<li><code>jax.lib.xla_bridge.xla_client</code>: use
<code>jax.lib.xla_client</code> directly.</li>
<li><code>jax.lib.xla_bridge.get_backend</code>: use
<code>jax.extend.backend.get_backend</code>.</li>
<li><code>jax.lib.xla_bridge.default_backend</code>: use
<code>jax.extend.backend.default_backend</code>.</li>
</ul>
</li>
<li>The <code>jax.experimental.array_api</code> module is deprecated,
and importing it is no
longer required to use the Array API. <code>jax.numpy</code> supports
the array API
directly; see <code>python-array-api</code> for more information.</li>
<li>The internal utilities <code>jax.core.check_eqn</code>,
<code>jax.core.check_type</code>, and
<code>jax.core.check_valid_jaxtype</code> are now deprecated, and will
be removed in
the future.</li>
<li><code>jax.numpy.round_</code> has been deprecated, following removal
of the corresponding</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/google/jax/blob/main/CHANGELOG.md">jax's
changelog</a>.</em></p>
<blockquote>
<h2>jax 0.4.32 (September 11, 2024)</h2>
<ul>
<li>
<p>New Functionality</p>
<ul>
<li>Added {func}<code>jax.extend.ffi.ffi_call</code> and
{func}<code>jax.extend.ffi.ffi_lowering</code>
to support the use of the new {ref}<code>ffi-tutorial</code> to
interface with custom
C++ and CUDA code from JAX.</li>
</ul>
</li>
<li>
<p>Changes</p>
<ul>
<li><code>jax_pmap_no_rank_reduction</code> flag is set to
<code>True</code> by default.
<ul>
<li>array[0] on a pmap result now introduces a reshape (use array[0:1]
instead).</li>
<li>The per-shard shape (accessable via jax_array.addressable_shards or
jax_array.addressable_data(0)) now has a leading (1, ...). Update code
that directly accesses shards accordingly. The rank of the
per-shard-shape
now matches that of the global shape which is the same behavior as jit.
This avoids costly reshapes when passing results from pmap into
jit.</li>
</ul>
</li>
<li><code>jax_enable_memories</code> flag is set to <code>True</code> by
default.</li>
<li>{mod}<code>jax.numpy</code> now supports v2023.12 of the Python
Array API Standard.
See {ref}<code>python-array-api</code> for more information.</li>
<li>Computations on the CPU backend may now be dispatched asynchronously
in
more cases. Previously non-parallel computations were always dispatched
synchronously. You can recover the old behavior by setting
<code>jax.config.update('jax_cpu_enable_async_dispatch',
False)</code>.</li>
<li>Added new {func}<code>jax.process_indices</code> function to replace
the
<code>jax.host_ids()</code> function that was deprecated in JAX
v0.2.13.</li>
<li>To align with the behavior of <code>numpy.fabs</code>,
<code>jax.numpy.fabs</code> has been
modified to no longer support <code>complex dtypes</code>.</li>
<li><code>jax.tree_util.register_dataclass</code> now checks that
<code>data_fields</code>
and <code>meta_fields</code> includes all dataclass fields with
<code>init=True</code>
and only them, if <code>nodetype</code> is a dataclass.</li>
<li>Several {mod}<code>jax.numpy</code> functions now have full
{class}<code>~jax.numpy.ufunc</code>
interfaces, including {obj}<code>~jax.numpy.add</code>,
{obj}<code>~jax.numpy.multiply</code>,
{obj}<code>~jax.numpy.bitwise_and</code>,
{obj}<code>~jax.numpy.bitwise_or</code>,
{obj}<code>~jax.numpy.bitwise_xor</code>,
{obj}<code>~jax.numpy.logical_and</code>,
{obj}<code>~jax.numpy.logical_and</code>, and
{obj}<code>~jax.numpy.logical_and</code>.
In future releases we plan to expand these to other ufuncs.</li>
<li>Added {func}<code>jax.lax.optimization_barrier</code>, which allows
users to prevent
compiler optimizations such as common-subexpression elimination and to
control scheduling.</li>
</ul>
</li>
<li>
<p>Breaking changes</p>
<ul>
<li>The MHLO MLIR dialect (<code>jax.extend.mlir.mhlo</code>) has been
removed. Use the
<code>stablehlo</code> dialect instead.</li>
</ul>
</li>
<li>
<p>Deprecations</p>
<ul>
<li>Complex inputs to {func}<code>jax.numpy.clip</code> and
{func}<code>jax.numpy.hypot</code> are
no longer allowed, after being deprecated since JAX v0.4.27.</li>
<li>Deprecated the following APIs:
<ul>
<li><code>jax.lib.xla_bridge.xla_client</code>: use
{mod}<code>jax.lib.xla_client</code> directly.</li>
<li><code>jax.lib.xla_bridge.get_backend</code>: use
{func}<code>jax.extend.backend.get_backend</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/google/jax/commit/1594d2f30fdbfebf693aba4a2b264e4a3e52acc6"><code>1594d2f</code></a>
Prepare for v0.4.32 release.</li>
<li><a
href="https://github.com/google/jax/commit/ed849ff9e0576dcee2514741b5ffa951a94e20a8"><code>ed849ff</code></a>
Make sure to call the superclass' <strong>init</strong>() on a newly
created instance in P...</li>
<li><a
href="https://github.com/google/jax/commit/2bd1fdead81581db08ee84a0d1f82c407ccd6b11"><code>2bd1fde</code></a>
Relax test tolerance in pinv test to fix a CI failure on Windows
CPU.</li>
<li><a
href="https://github.com/google/jax/commit/e869a9d65e568e36e95940db302f94f9b7b973c4"><code>e869a9d</code></a>
Merge pull request <a
href="https://github.com/google/jax/issues/23415">#23415</a>
from kaixih:key_value_seq_lengths</li>
<li><a
href="https://github.com/google/jax/commit/ea68f4569c5474f20e52b96ab88c287ab843130a"><code>ea68f45</code></a>
Internal change</li>
<li><a
href="https://github.com/google/jax/commit/49dd6ed8d891ee6b7bbfcf7cc425382a7235556b"><code>49dd6ed</code></a>
Disable a pallas export compatibility test that fails on TPU v6e.</li>
<li><a
href="https://github.com/google/jax/commit/808003b4e29e878349192e0f63fa1a2454ace56b"><code>808003b</code></a>
Update users of jax.tree.map() to be more careful about how they handle
Nones.</li>
<li><a
href="https://github.com/google/jax/commit/e3c4b20fa04893ad986c3184387fbd3817f1515d"><code>e3c4b20</code></a>
[Pallas] Implement tiled and swizzled Memref loads for Mosaic GPU via
&quot;GPUBlo...</li>
<li><a
href="https://github.com/google/jax/commit/c659dc9a011bf8ff604a7e23f916920ff717288b"><code>c659dc9</code></a>
[Pallas] Disable win32 gpu_ops_test.</li>
<li><a
href="https://github.com/google/jax/commit/14b86259d5a7fc31235f3fdfb3bcb64aadab338b"><code>14b8625</code></a>
Merge pull request <a
href="https://github.com/google/jax/issues/23549">#23549</a>
from pschuh:docs-update</li>
<li>Additional commits viewable in <a
href="https://github.com/google/jax/compare/jax-v0.4.31...jax-v0.4.32">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.31&new-version=0.4.32)](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>

pyproject.toml

@@ -34,7 +34,7 @@ dev = [
     "pyright==1.1.345",
 ]
 gui = ["textual==0.79.1", "pyclip==0.7"]
-jax = ["jax==0.4.31", "numpy==2.1.1"]
+jax = ["jax==0.4.32", "numpy==2.1.1"]
 onnx = ["onnx==1.16.2", "numpy==2.1.1"]
 riscv = ["riscemu==2.2.7"]
 wgpu = ["wgpu==0.16.0"]