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

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX release v0.4.24

Changelog

Sourced from jax's changelog.

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. This is needed because custom_partitioning and JAX callbacks need physical devices to create Shardings during lowering. This is a temporary state until we can create Shardings without physical devices.
  • {func}jax.numpy.argsort and {func}jax.numpy.sort now support the stable and descending arguments.
  • Several changes to the handling of shape polymorphism (used in {mod}jax.experimental.jax2tf and {mod}jax.experimental.export):
    • cleaner pretty-printing of symbolic expressions ({jax-issue}[#19227](https://github.com/google/jax/issues/19227))
    • added the ability to specify symbolic constraints on the dimension variables. This makes shape polymorphism more expressive, and gives a way to workaround limitations in the reasoning about inequalities. See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#user-specified-symbolic-constraints.
    • with the addition of symbolic constraints ({jax-issue}[#19235](https://github.com/google/jax/issues/19235)) we now consider dimension variables from different scopes to be different, even if they have the same name. Symbolic expressions from different scopes cannot interact, e.g., in arithmetic operations. Scopes are introduced by {func}jax.experimental.jax2tf.convert, {func}jax.experimental.export.symbolic_shape, {func}jax.experimental.export.symbolic_args_specs. The scope of a symbolic expression e can be read with e.scope and passed in to the above functions to direct them to construct sybolic expressions in a given scope. See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#user-specified-symbolic-constraints.
    • simplified and faster equality comparisons, where we consider two symbolic dimensions to be equal if the normalized form of their difference reduces to 0 ({jax-issue}[#19231](https://github.com/google/jax/issues/19231); note that this may result in user-visible behavior changes)
    • improved the error messages for inconclusive inequality comparisons ({jax-issue}[#19235](https://github.com/google/jax/issues/19235)).
    • the core.non_negative_dim API (introduced recently) was deprecated and core.max_dim and core.min_dim were introduced ({jax-issue}[#18953](https://github.com/google/jax/issues/18953)) to express max and min for symbolic dimensions. You can use core.max_dim(d, 0) instead of core.non_negative_dim(d).
    • the shape_poly.is_poly_dim is deprecated in favor of export.is_symbolic_dim ({jax-issue}[#19282](https://github.com/google/jax/issues/19282)).
    • the export.args_specs is deprecated in favor of export.symbolic_args_specs ({jax-issue}#19283`).
    • the shape_poly.PolyShape and jax2tf.PolyShape are deprecated, use strings for polymorphic shapes specifications ({jax-issue}[#19284](https://github.com/google/jax/issues/19284)).
    • JAX default native serialization version is now 9. This is relevant for {mod}jax.experimental.jax2tf and {mod}jax.experimental.export. See description of version numbers.

... (truncated)

Commits

• 9e94e6e Fixed a typo in min/max Triton lowering rules
• b53f757 Merge pull request #19667 from jakevdp:array-empty-repr
• c1c0c1c Merge pull request #19634 from jakevdp:key-reuse-scan
• d9cbd7b Improve repr for empty jax.Array
• 69a9f7f Merge pull request #19629 from jakevdp:key-reuse-pjit
• 206398a Merge pull request #19599 from ROCm:rocm-add-triton_command_buffer
• be99451 Merge pull request #19665 from jakevdp:disable-jit-doc
• 82611eb document that under disable_jit, individual primitives are still compiled
• f01c27f [ROCm]: Add ROCm command buffer support for triton kernel
• e224c3d Update XLA dependency to use revision
• 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]

Superseded by #136.