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

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.5.3

• New Features

  • Added a allow_negative_indices option to jax.lax.dynamic_slice, jax.lax.dynamic_update_slice and related functions. The default is true, matching the current behavior. If set to false, JAX does not need to emit code clamping negative indices, which improves code size.
  • Added a replace option to jax.random.categorical to enable sampling without replacement.

Changelog

Sourced from jax's changelog.

jax 0.5.3 (Mar 19, 2025)

• New Features

  • Added a allow_negative_indices option to {func}jax.lax.dynamic_slice, {func}jax.lax.dynamic_update_slice and related functions. The default is true, matching the current behavior. If set to false, JAX does not need to emit code clamping negative indices, which improves code size.
  • Added a replace option to {func}jax.random.categorical to enable sampling without replacement.

jax 0.5.2 (Mar 4, 2025)

Patch release of 0.5.1

• Bug fixes
  • Fixes TPU metric logging and tpu-info, which was broken in 0.5.1

jax 0.5.1 (Feb 24, 2025)

• New Features

  • Added an experimental {func}jax.experimental.custom_dce.custom_dce decorator to support customizing the behavior of opaque functions under JAX-level dead code elimination (DCE). See {jax-issue}[#25956](https://github.com/jax-ml/jax/issues/25956) for more details.
  • Added low-level reduction APIs in {mod}jax.lax: {func}jax.lax.reduce_sum, {func}jax.lax.reduce_prod, {func}jax.lax.reduce_max, {func}jax.lax.reduce_min, {func}jax.lax.reduce_and, {func}jax.lax.reduce_or, and {func}jax.lax.reduce_xor.
  • {func}jax.lax.linalg.qr, and {func}jax.scipy.linalg.qr, now support column-pivoting on CPU and GPU. See {jax-issue}[#20282](https://github.com/jax-ml/jax/issues/20282) and
  • Added {func}jax.random.multinomial. {jax-issue}[#25955](https://github.com/jax-ml/jax/issues/25955) for more details.

• Changes

  • JAX_CPU_COLLECTIVES_IMPLEMENTATION and JAX_NUM_CPU_DEVICES now work as env vars. Before they could only be specified via jax.config or flags.
  • JAX_CPU_COLLECTIVES_IMPLEMENTATION now defaults to 'gloo', meaning multi-process CPU communication works out-of-the-box.
  • The jax[tpu] TPU extra no longer depends on the libtpu-nightly package. This package may safely be removed if it is present on your machine; JAX now uses libtpu instead.

• Deprecations

  • The internal function linear_util.wrap_init and the constructor core.Jaxpr now must take a non-empty core.DebugInfo kwarg. For a limited time, a DeprecationWarning is printed if jax.extend.linear_util.wrap_init is used without debugging info. A downstream effect of this several other internal functions need debug info. This change does not affect public APIs. See jax-ml/jax#26480 for more detail.

... (truncated)

Commits

• c8032a9 Fix line continuation character in Windows wheel build.
• ed43119 JAX release v0.5.3
• 8c7a55e Update XLA dependency to use revision
• 663ef7a Check the type of mesh in use_abstract_mesh and use_concrete_mesh
• 3f91b4b Move jaxlib/{cuda,rocm}_plugin_extension into jaxlib/{cuda/rocm}/
• 01a110c Better mosaic lowering for dynamic shapes, extend an interpreter into shape_p...
• 0fb5974 Support tuples in custom_partitioning.
• 080804c Fix logging_test fails on Linux with NVIDIA Driver only.
• 54691b1 [Mosaic GPU] Support reads/writes from SMEM to WGMMARowFragLayout arrays.
• 76d9890 Run the stream annotation tests on 2 devices so that it can be tested in TAP
• 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)

[dependabot]

Assignees

The following users could not be added as assignees: akuporos, jiwaszki, p-wysocki, rkazants. Either they do not exist or they do not have the correct permissions to be added as an assignee.

Please fix the above issues or remove invalid values from dependabot.yml.