Update jax requirement from <=0.4.14 to <=0.4.25 in /tests

[dependabot]

Updates the requirements on jax to permit the latest version.

Changelog

Sourced from jax's changelog.

jax 0.4.25 (Feb 26, 2024)

• New Features

  • Added CUDA Array Interface import support (requires jaxlib 0.4.24).
  • JAX arrays now support NumPy-style scalar boolean indexing, e.g. x[True] or x[False].
  • Added {mod}jax.tree module, with a more convenient interface for referencing functions in {mod}jax.tree_util.
  • {func}jax.tree.transpose (i.e. {func}jax.tree_util.tree_transpose) now accepts inner_treedef=None, in which case the inner treedef will be automatically inferred.

• Changes

  • Pallas now uses XLA instead of the Triton Python APIs to compile Triton kernels. You can revert to the old behavior by setting the JAX_TRITON_COMPILE_VIA_XLA environment variable to "0".
  • Several deprecated APIs in {mod}jax.interpreters.xla that were removed in v0.4.24 have been re-added in v0.4.25, including backend_specific_translations, translations, register_translation, xla_destructure, TranslationRule, TranslationContext, and XLAOp. These are still considered deprecated, and will be removed again in the future when better replacements are available. Refer to {jax-issue}[#19816](https://github.com/google/jax/issues/19816) for discussion.

• Deprecations & Removals

  • {func}jax.numpy.linalg.solve now shows a deprecation warning for batched 1D solves with b.ndim > 1. In the future these will be treated as batched 2D solves.
  • Conversion of a non-scalar array to a Python scalar now raises an error, regardless of the size of the array. Previously a deprecation warning was raised in the case of non-scalar arrays of size 1. This follows a similar deprecation in NumPy.
  • The previously deprecated configuration APIs have been removed following a standard 3 months deprecation cycle (see {ref}api-compatibility). These include
    • the jax.config.config object and
    • the define_*_state and DEFINE_* methods of {data}jax.config.
  • Importing the jax.config submodule via import jax.config is deprecated. To configure JAX use import jax and then reference the config object via jax.config.
  • The minimum jaxlib version is now 0.4.20.

jaxlib 0.4.25 (Feb 26, 2024)

jax 0.4.24 (Feb 6, 2024)

• Changes

  • JAX lowering to StableHLO does not depend on physical devices anymore. If your primitive wraps custom_paritioning or JAX callbacks in the lowering rule i.e. function passed to rule parameter of mlir.register_lowering then add your primitive to jax._src.dispatch.prim_requires_devices_during_lowering set.

... (truncated)

Commits

• f9e20d5 Start the release of jax and jaxlib 0.4.25
• f4045dc Remove the deprecation of jax.tree_map for the release of 0.4.25
• 22996bb Update XLA dependency to use revision
• 61aa7e8 [Mosaic] Fix bug in divisibility check in infer_vector_layout load and store ...
• 072b43b Merge pull request #19960 from mattjj:attrs-autodiff
• b0b88d8 [attrs] add linearize and vjp support
• ae9d677 Merge pull request #19936 from tillahoffmann:fix-19935
• e1e9de0 Merge pull request #19499 from nstarman:array_api-broadcast_to-type-hint
• 75cdef7 [Mosaic][NFC] Prefer mlir aliases for llvm classes/functions within mlir name...
• 8a43140 [Mosaic][apply_vector_layout][NFC] Use LLVM_UNLIKELY in TPU_ASSERT_* macros
• 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.