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

[dependabot]

Updates the requirements on jax to permit the latest version.

Changelog

Sourced from jax's changelog.

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.

jaxlib 0.4.17 (Oct 3, 2023)

• Changes:

  • Python 3.12 wheels were added in this release.
  • The CUDA 12 wheels now require CUDA 12.2 or newer and cuDNN 8.9.4 or newer.

• Bug fixes:

  • Fixed log spam from ABSL when the JAX CPU backend was initialized.

jax 0.4.16 (Sept 18, 2023)

• Changes
  • Added {class}jax.numpy.ufunc, as well as {func}jax.numpy.frompyfunc, which can convert any scalar-valued function into a {func}numpy.ufunc-like object, with methods such as {meth}~jax.numpy.ufunc.outer, {meth}~jax.numpy.ufunc.reduce, {meth}~jax.numpy.ufunc.accumulate, {meth}~jax.numpy.ufunc.at, and {meth}~jax.numpy.ufunc.reduceat ({jax-issue}[#17054](https://github.com/google/jax/issues/17054)).
  • Added {func}jax.scipy.integrate.trapezoid.
  • When not running under IPython: when an exception is raised, JAX now filters out the entirety of its internal frames from tracebacks. (Without the "unfiltered stack trace" that previously appeared.) This should produce much friendlier-looking tracebacks. See here for an example. This behavior can be changed by setting JAX_TRACEBACK_FILTERING=remove_frames (for two

... (truncated)

Commits

• 4cb5eee Merge pull request #17905 from skye:version
• b9c602b Update XLA commit and versions for jax 0.4.17 release
• 3d84808 Merge pull request #17851 from mattjj:readme-tweak
• dddbe43 Update README.md
• 24ad445 [Pallas] Add support for pytrees in scalar prefetch
• ee8af09 Fix mock_gpu_test on OSS build.
• c3e73c6 Merge pull request #17760 from superbobry:array-any
• b2ac2de [XLA] Split --xla_detailed_logging_and_dumping debug flag into --xla_detailed...
• 72407f6 Update XLA dependency to use revision
• 17e259b fix typo: device_fun(c) -> device_fun(x)
• 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 #147.