pip prod(deps): bump jax from 0.4.34 to 0.4.35 (#3332)

Bumps [jax](https://github.com/jax-ml/jax) from 0.4.34 to 0.4.35.
<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.35</h2>
<ul>
<li>
<p>Breaking Changes</p>
<ul>
<li><code>jax.numpy.isscalar</code> now returns True for any array-like
object with
zero dimensions. Previously it only returned True for zero-dimensional
array-like objects with a weak dtype.</li>
<li><code>jax.experimental.host_callback</code> has been deprecated
since March 2024, with
JAX version 0.4.26. Now we removed it.
See <code>[#20385](https://github.com/jax-ml/jax/issues/20385)</code>
for a discussion of alternatives.</li>
</ul>
</li>
<li>
<p>Changes:</p>
<ul>
<li><code>jax.lax.FftType</code> was introduced as a public name for the
enum of FFT
operations. The semi-public API <code>jax.lib.xla_client.FftType</code>
has been
deprecated.</li>
<li>TPU: JAX now installs TPU support from the <code>libtpu</code>
package rather than
<code>libtpu-nightly</code>. For the next few releases JAX will pin an
empty version of
<code>libtpu-nightly</code> as well as <code>libtpu</code> to ease the
transition; that dependency
will be removed in Q1 2025.</li>
</ul>
</li>
<li>
<p>Deprecations:</p>
<ul>
<li>The semi-public API <code>jax.lib.xla_client.PaddingType</code> has
been deprecated.
No JAX APIs consume this type, so there is no replacement.</li>
<li>The default behavior of <code>jax.pure_callback</code> and
<code>jax.extend.ffi.ffi_call</code> under <code>vmap</code> has been
deprecated and so has
the <code>vectorized</code> parameter to those functions. The
<code>vmap_method</code> parameter
should be used instead for better defined behavior. See the discussion
in
<code>[#23881](https://github.com/jax-ml/jax/issues/23881)</code> for
more details.</li>
<li>The semi-public API
<code>jax.lib.xla_client.register_custom_call_target</code> has
been deprecated. Use the JAX FFI instead.</li>
<li>The semi-public APIs <code>jax.lib.xla_client.dtype_to_etype</code>,
<code>jax.lib.xla_client.ops</code>,
<code>jax.lib.xla_client.shape_from_pyval</code>,
<code>jax.lib.xla_client.PrimitiveType</code>,
<code>jax.lib.xla_client.Shape</code>,
<code>jax.lib.xla_client.XlaBuilder</code>, and
<code>jax.lib.xla_client.XlaComputation</code> have been deprecated. Use
StableHLO
instead.</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.35</h2>
<ul>
<li>
<p>Breaking Changes</p>
<ul>
<li>{func}<code>jax.numpy.isscalar</code> now returns True for any
array-like object with
zero dimensions. Previously it only returned True for zero-dimensional
array-like objects with a weak dtype.</li>
<li><code>jax.experimental.host_callback</code> has been deprecated
since March 2024, with
JAX version 0.4.26. Now we removed it.
See
{jax-issue}<code>[#20385](https://github.com/jax-ml/jax/issues/20385)</code>
for a discussion of alternatives.</li>
</ul>
</li>
<li>
<p>Changes:</p>
<ul>
<li><code>jax.lax.FftType</code> was introduced as a public name for the
enum of FFT
operations. The semi-public API <code>jax.lib.xla_client.FftType</code>
has been
deprecated.</li>
<li>TPU: JAX now installs TPU support from the <code>libtpu</code>
package rather than
<code>libtpu-nightly</code>. For the next few releases JAX will pin an
empty version of
<code>libtpu-nightly</code> as well as <code>libtpu</code> to ease the
transition; that dependency
will be removed in Q1 2025.</li>
</ul>
</li>
<li>
<p>Deprecations:</p>
<ul>
<li>The semi-public API <code>jax.lib.xla_client.PaddingType</code> has
been deprecated.
No JAX APIs consume this type, so there is no replacement.</li>
<li>The default behavior of {func}<code>jax.pure_callback</code> and
{func}<code>jax.extend.ffi.ffi_call</code> under <code>vmap</code> has
been deprecated and so has
the <code>vectorized</code> parameter to those functions. The
<code>vmap_method</code> parameter
should be used instead for better defined behavior. See the discussion
in

{jax-issue}<code>[#23881](https://github.com/jax-ml/jax/issues/23881)</code>
for more details.</li>
<li>The semi-public API
<code>jax.lib.xla_client.register_custom_call_target</code> has
been deprecated. Use the JAX FFI instead.</li>
<li>The semi-public APIs <code>jax.lib.xla_client.dtype_to_etype</code>,
<code>jax.lib.xla_client.ops</code>,
<code>jax.lib.xla_client.shape_from_pyval</code>,
<code>jax.lib.xla_client.PrimitiveType</code>,
<code>jax.lib.xla_client.Shape</code>,
<code>jax.lib.xla_client.XlaBuilder</code>, and
<code>jax.lib.xla_client.XlaComputation</code> have been deprecated. Use
StableHLO
instead.</li>
</ul>
</li>
</ul>
</blockquote>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a
href="https://github.com/jax-ml/jax/commit/81991d87c8608d07f9752940375092aa4fb9610e"><code>81991d8</code></a>
JAX release 0.4.35</li>
<li><a
href="https://github.com/jax-ml/jax/commit/e4f3f8f064228c0739b3333a9e78ac715cb34273"><code>e4f3f8f</code></a>
Use libtpu releases rather than libtpu-nightly for jax[tpu].</li>
<li><a
href="https://github.com/jax-ml/jax/commit/1c6b0a9193befffb1a074382ffc8ce0ae86d4900"><code>1c6b0a9</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/24465">#24465</a>
from jakevdp:fix-mypy</li>
<li><a
href="https://github.com/jax-ml/jax/commit/9a2dd19a92942a8a00a59d314cf6fd534eea9be3"><code>9a2dd19</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/21524">#21524</a>
from andportnoy:aportnoy/unknown-platform-lowering-...</li>
<li><a
href="https://github.com/jax-ml/jax/commit/1e41d5ef6f20c0018aa1ab69e021804535a9229d"><code>1e41d5e</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/24452">#24452</a>
from jakevdp:insert-doc</li>
<li><a
href="https://github.com/jax-ml/jax/commit/1a2737b72b21b4341ff573660048e773c6bc52ad"><code>1a2737b</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/24467">#24467</a>
from andportnoy:patch-2</li>
<li><a
href="https://github.com/jax-ml/jax/commit/849850216d4e0df9205f6b269f6c2ac855ac5658"><code>8498502</code></a>
fix mypy error</li>
<li><a
href="https://github.com/jax-ml/jax/commit/92aa9a399013924ccb799d4d987520cb14d20928"><code>92aa9a3</code></a>
Merge pull request <a
href="https://redirect.github.com/jax-ml/jax/issues/24463">#24463</a>
from jakevdp:ml-dtypes-version</li>
<li><a
href="https://github.com/jax-ml/jax/commit/c60bafcc335967e7d1a35b9a7af1e23c0d593531"><code>c60bafc</code></a>
[Pallas TPU] Fix lowering for <code>jnp.remainder</code></li>
<li><a
href="https://github.com/jax-ml/jax/commit/637898493ead441a82cac3704f3f274ad973af8d"><code>6378984</code></a>
Add back the import of <code>jtu</code> in
<code>flash_attention.py</code></li>
<li>Additional commits viewable in <a
href="https://github.com/jax-ml/jax/compare/jax-v0.4.34...jax-v0.4.35">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.34&new-version=0.4.35)](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>

pyproject.toml

@@ -34,7 +34,7 @@ dev = [
     "pyright==1.1.385",
 ]
 gui = ["textual==0.83.0", "pyclip==0.7"]
-jax = ["jax==0.4.34", "numpy==2.1.2"]
+jax = ["jax==0.4.35", "numpy==2.1.2"]
 onnx = ["onnx==1.17.0", "numpy==2.1.2"]
 riscv = ["riscemu==2.2.7"]
 wgpu = ["wgpu==0.18.1"]