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

[dependabot]

Updates the requirements on jax to permit the latest version.

Changelog

Sourced from jax's changelog.

jax 0.4.18 (Oct 6, 2023)

jaxlib 0.4.18 (Oct 6, 2023)

• Changes

  • CUDA jaxlibs now depend on the user to install a compatible NCCL version. If using the recommended cuda12_pip installation, NCCL should be installed automatically. Currently, NCCL 2.16 or newer is required.
  • We now provide Linux aarch64 wheels, both with and without NVIDIA GPU support.

• Deprecations

  • A number of internal utilities and inadvertent exports in {mod}jax.lax have been deprecated, and will be removed in a future release.
    • jax.lax.dtypes: use jax.dtypes instead.
    • jax.lax.itertools: use itertools instead.
    • naryop, naryop_dtype_rule, standard_abstract_eval, standard_naryop, standard_primitive, standard_unop, unop, and unop_dtype_rule are internal utilities, now deprecated without replacement.

• Bug fixes

  • Fixed Cloud TPU regression where compilation would OOM due to smem.

jax 0.4.17 (Oct 3, 2023)

• New features
  • Added new {func}jax.numpy.bitwise_count function, matching the API of the simlar function recently added to NumPy.
• Deprecations
  • Removed the deprecated module jax.abstract_arrays and all its contents.
  • Named key constructors in {mod}jax.random are deprecated. Pass the impl argument to {func}jax.random.PRNGKey or {func}jax.random.key instead:
    • random.threefry2x32_key(seed) becomes random.PRNGKey(seed, impl='threefry2x32')
    • random.rbg_key(seed) becomes random.PRNGKey(seed, impl='rbg')
    • random.unsafe_rbg_key(seed) becomes random.PRNGKey(seed, impl='unsafe_rbg')
• Changes:
  • CUDA: JAX now verifies that the CUDA libraries it finds are at least as new as the CUDA libraries that JAX was built against. If older libraries are found, JAX raises an exception since that is preferable to mysterious failures and crashes.
  • Removed the "No GPU/TPU" found warning. Instead warn if, on Linux, an NVIDIA GPU or a Google TPU are found but not used and --jax_platforms was not specified.
  • {func}jax.scipy.stats.mode now returns a 0 count if the mode is taken across a size-0 axis, matching the behavior of scipy.stats.mode in SciPy 1.11.
  • Most jax.numpy functions and attributes now have fully-defined type stubs. Previously many of these were treated as Any by static type checkers like mypy and pytype.

... (truncated)

Commits

• e27af48 Merge pull request #17989 from skye:version
• d4a1bb9 Update setup.py and CHANGELOG for jax 0.4.18 release
• f0e4ea2 Merge pull request #17987 from jakevdp:lax-dep
• 6681e64 Merge pull request #17962 from jakevdp:fix-bitwise-count
• fae53d9 Merge pull request #17959 from jakevdp:lax-abs
• 5ac2513 Merge pull request #17988 from froystig:rng-spec-jex
• 31b680b Merge pull request #17981 from hawkinsp:docs2
• 8c4d020 Improve CUDA install documentation.
• ce6a0c4 jax.lax: deprecate inadvertent exports & internal utilities
• 7b831ba test custom PRNG impl construction round trip
• 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 #153.