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

[dependabot]

Updates the requirements on jax to permit the latest version.

Release notes

Sourced from jax's releases.

JAX v0.4.28

• Bug fixes

  • Reverted a change to make_jaxpr that was breaking Equinox (#21116).

• Deprecations & removals

  • The kind argument to jax.numpy.sort and jax.numpy.argsort is now removed. Use stable=True or stable=False instead.
  • Removed get_compute_capability from the jax.experimental.pallas.gpu module. Use the compute_capability attribute of a GPU device, returned by jax.devices or jax.local_devices, instead.

• Changes

  • The minimum jaxlib version of this release is 0.4.27.

Changelog

Sourced from jax's changelog.

jax 0.4.28

• Bug fixes

  • Reverted a change to make_jaxpr that was breaking Equinox (#21116).

• Deprecations & removals

  • The kind argument to {func}jax.numpy.sort and {func}jax.numpy.argsort is now removed. Use stable=True or stable=False instead.
  • Removed get_compute_capability from the jax.experimental.pallas.gpu module. Use the compute_capability attribute of a GPU device, returned by {func}jax.devices or {func}jax.local_devices, instead.
  • The newshape argument to {func}jax.numpy.reshapeis being deprecated and will soon be removed. Use shape instead.

• Changes

  • The minimum jaxlib version of this release is 0.4.27.

jaxlib 0.4.28

• Bug fixes

  • Fixes a memory corruption bug in the type name of Array and JIT Python objects in Python 3.10 or earlier.
  • Fixed a warning '+ptx84' is not a recognized feature for this target under CUDA 12.4.
  • Fixed a slow compilation problem on CPU.

• Changes

  • The Windows build is now built with Clang instead of MSVC.

jax 0.4.27 (May 7, 2024)

• New Functionality

  • Added {func}jax.numpy.unstack and {func}jax.numpy.cumulative_sum, following their addition in the array API 2023 standard, soon to be adopted by NumPy.
  • Added a new config option jax_cpu_collectives_implementation to select the implementation of cross-process collective operations used by the CPU backend. Choices available are 'none'(default), 'gloo' and 'mpi' (requires jaxlib 0.4.26). If set to 'none', cross-process collective operations are disabled.

• Changes

  • {func}jax.pure_callback, {func}jax.experimental.io_callback and {func}jax.debug.callback now use {class}jax.Array instead of {class}np.ndarray. You can recover the old behavior by transforming the arguments via jax.tree.map(np.asarray, args) before passing them to the callback.
  • complex_arr.astype(bool) now follows the same semantics as NumPy, returning False where complex_arr is equal to 0 + 0j, and True otherwise.
  • core.Token now is a non-trivial class which wraps a jax.Array. It could

... (truncated)

Commits

• 9e62994 Merge pull request #21135 from hawkinsp:release
• 038dfee Prepare 0.4.28 release.
• f98e707 Update XLA dependency to use revision
• 1a7a2aa Merge pull request #21106 from jakevdp:linalg-precision
• 0c4d81c Merge pull request #21138 from jakevdp:einsum-doc
• 2ddb7ff jnp.linalg: add precision & preferred_element_type to dot-like functions
• eb0b1b0 Merge pull request #21108 from justinjfu/skip_pallas_test_64
• 671fb12 Update the multi-process note in pjit's docstring
• 2be3f6d Merge pull request #21146 from jakevdp:fix-multidot
• 89d25bb Reenable examples_test in Bazel build.
• 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.