Update jax requirement from <=0.4.31 to <=0.4.35 in /tests

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.4.35

• Breaking Changes

  • jax.numpy.isscalar 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.
  • jax.experimental.host_callback has been deprecated since March 2024, with JAX version 0.4.26. Now we removed it. See [#20385](https://github.com/jax-ml/jax/issues/20385) for a discussion of alternatives.

• Changes:

  • jax.lax.FftType was introduced as a public name for the enum of FFT operations. The semi-public API jax.lib.xla_client.FftType has been deprecated.
  • TPU: JAX now installs TPU support from the libtpu package rather than libtpu-nightly. For the next few releases JAX will pin an empty version of libtpu-nightly as well as libtpu to ease the transition; that dependency will be removed in Q1 2025.

• Deprecations:

  • The semi-public API jax.lib.xla_client.PaddingType has been deprecated. No JAX APIs consume this type, so there is no replacement.
  • The default behavior of jax.pure_callback and jax.extend.ffi.ffi_call under vmap has been deprecated and so has the vectorized parameter to those functions. The vmap_method parameter should be used instead for better defined behavior. See the discussion in [#23881](https://github.com/jax-ml/jax/issues/23881) for more details.
  • The semi-public API jax.lib.xla_client.register_custom_call_target has been deprecated. Use the JAX FFI instead.
  • The semi-public APIs jax.lib.xla_client.dtype_to_etype, jax.lib.xla_client.ops, jax.lib.xla_client.shape_from_pyval, jax.lib.xla_client.PrimitiveType, jax.lib.xla_client.Shape, jax.lib.xla_client.XlaBuilder, and jax.lib.xla_client.XlaComputation have been deprecated. Use StableHLO instead.

Changelog

Sourced from jax's changelog.

jax 0.4.35 (Oct 22, 2024)

• Breaking Changes

  • {func}jax.numpy.isscalar 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.
  • jax.experimental.host_callback has been deprecated since March 2024, with JAX version 0.4.26. Now we removed it. See {jax-issue}[#20385](https://github.com/jax-ml/jax/issues/20385) for a discussion of alternatives.

• Changes:

  • jax.lax.FftType was introduced as a public name for the enum of FFT operations. The semi-public API jax.lib.xla_client.FftType has been deprecated.
  • TPU: JAX now installs TPU support from the libtpu package rather than libtpu-nightly. For the next few releases JAX will pin an empty version of libtpu-nightly as well as libtpu to ease the transition; that dependency will be removed in Q1 2025.

• Deprecations:

  • The semi-public API jax.lib.xla_client.PaddingType has been deprecated. No JAX APIs consume this type, so there is no replacement.
  • The default behavior of {func}jax.pure_callback and {func}jax.extend.ffi.ffi_call under vmap has been deprecated and so has the vectorized parameter to those functions. The vmap_method parameter should be used instead for better defined behavior. See the discussion in {jax-issue}[#23881](https://github.com/jax-ml/jax/issues/23881) for more details.
  • The semi-public API jax.lib.xla_client.register_custom_call_target has been deprecated. Use the JAX FFI instead.
  • The semi-public APIs jax.lib.xla_client.dtype_to_etype, jax.lib.xla_client.ops, jax.lib.xla_client.shape_from_pyval, jax.lib.xla_client.PrimitiveType, jax.lib.xla_client.Shape, jax.lib.xla_client.XlaBuilder, and jax.lib.xla_client.XlaComputation have been deprecated. Use StableHLO instead.

jax 0.4.34 (October 4, 2024)

• New Functionality

  • This release includes wheels for Python 3.13. Free-threading mode is not yet supported.
  • jax.errors.JaxRuntimeError has been added as a public alias for the formerly private XlaRuntimeError type.

• Breaking changes

  • jax_pmap_no_rank_reduction flag is set to True by default.
    • array[0] on a pmap result now introduces a reshape (use array[0:1] instead).
    • The per-shard shape (accessable via jax_array.addressable_shards or jax_array.addressable_data(0)) now has a leading (1, ...). Update code

... (truncated)

Commits

• 81991d8 JAX release 0.4.35
• e4f3f8f Use libtpu releases rather than libtpu-nightly for jax[tpu].
• 1c6b0a9 Merge pull request #24465 from jakevdp:fix-mypy
• 9a2dd19 Merge pull request #21524 from andportnoy:aportnoy/unknown-platform-lowering-...
• 1e41d5e Merge pull request #24452 from jakevdp:insert-doc
• 1a2737b Merge pull request #24467 from andportnoy:patch-2
• 8498502 fix mypy error
• 92aa9a3 Merge pull request #24463 from jakevdp:ml-dtypes-version
• c60bafc [Pallas TPU] Fix lowering for jnp.remainder
• 6378984 Add back the import of jtu in flash_attention.py
• Additional commits viewable in compare view

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 commands and options

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)

[github-actions]

This PR will be closed in a week because of 2 weeks of no activity.

[github-actions]

This PR was closed because it has been stalled for 2 week with no activity.

[dependabot]

OK, I won't notify you again about this release, but will get in touch when a new version is available. If you'd rather skip all updates until the next major or minor version, let me know by commenting @dependabot ignore this major version or @dependabot ignore this minor version. You can also ignore all major, minor, or patch releases for a dependency by adding an ignore condition with the desired update_types to your config file.

If you change your mind, just re-open this PR and I'll resolve any conflicts on it.