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

[dependabot]

Updates the requirements on jax to permit the latest version.

Changelog

Sourced from jax's changelog.

jax 0.4.23 (Dec 13, 2023)

jaxlib 0.4.23 (Dec 13, 2023)

• Fixed a bug that caused verbose logging from the GPU compiler during compilation.

jax 0.4.22 (Dec 13, 2023)

• Deprecations
  • The device_buffer and device_buffers properties of JAX arrays are deprecated. Explicit buffers have been replaced by the more flexible array sharding interface, but the previous outputs can be recovered this way:
    • arr.device_buffer becomes arr.addressable_data(0)
    • arr.device_buffers becomes [x.data for x in arr.addressable_shards]

jaxlib 0.4.22 (Dec 13, 2023)

jax 0.4.21 (Dec 4 2023)

• New Features

  • Added {obj}jax.nn.squareplus.

• Changes

  • The minimum jaxlib version is now 0.4.19.
  • Released wheels are built now with clang instead of gcc.
  • Enforce that the device backend has not been initialized prior to calling jax.distributed.initialize().
  • Automate arguments to jax.distributed.initialize() in cloud TPU environments.

• Deprecations

  • The previously-deprecated sym_pos argument has been removed from {func}jax.scipy.linalg.solve. Use assume_a='pos' instead.
  • Passing None to {func}jax.array or {func}jax.asarray, either directly or within a list or tuple, is deprecated and now raises a {obj}FutureWarning. It currently is converted to NaN, and in the future will raise a {obj}TypeError.
  • Passing the condition, x, and y parameters to jax.numpy.where by keyword arguments has been deprecated, to match numpy.where.
  • Passing arguments to {func}jax.numpy.array_equal and {func}jax.numpy.array_equiv that cannot be converted to a JAX array is deprecated and now raises a {obj}DeprecationWaning. Currently the functions return False, in the future this will raise an exception.
  • The device() method of JAX arrays deprecated. Depending on the context, it may be replaced with one of the following:
    • {meth}jax.Array.devices returns the set of all devices used by the array.
    • {attr}jax.Array.sharding gives the sharding configuration used by the array.

jaxlib 0.4.21 (Dec 4 2023)

• Changes
  • In preparation for adding distributed CPU support, JAX now treats CPU

... (truncated)

Commits

• 1c68577 Prepare for emergency jax and jaxlib 0.4.23 release
• b392622 Add patch to suppress XLA:GPU logging.
• 25c16c0 Finish jax and jaxlib 0.4.22 release
• 3380b9f split the random generalized normal test and skip its K-S half
• 9198174 Merge pull request #18968 from 8bitmp3:update-thinking-in-jax
• 7edc80d Update thinking-in-jax working-with-pytrees
• 29ed3cd Merge pull request #18581 from 8bitmp3:jax-docs-thinking-in-jax
• b52bcc1 Reverts 3c07c10a9a55f9a32390dd10cf3f420bdf3f1ed8
• a4003bf Upgrade How to think in JAX
• e888806 Fix the pjit flakey test. The test was weirdly written in the first place. Th...
• 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.