⬆️ Update jax requirement from <0.4.16,>=0.4.1 to >=0.4.1,<0.4.17

[dependabot]

Updates the requirements on jax to permit the latest version.

Changelog

Sourced from jax's changelog.

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 separate unfiltered/filtered tracebacks, which was the old behavior) or JAX_TRACEBACK_FILTERING=off (for one unfiltered traceback).
  • jax2tf default serialization version is now 7, which introduces new shape safety assertions.
  • Devices passed to jax.sharding.Mesh should be hashable. This specifically applies to mock devices or user created devices. jax.devices() are already hashable.

• Breaking changes:

  • jax2tf now uses native serialization by default. See the jax2tf documentation for details and for mechanisms to override the default.
  • The option --jax_coordination_service has been removed. It is now always True.
  • jax.jaxpr_util has been removed from the public JAX namespace.
  • JAX_USE_PJRT_C_API_ON_TPU no longer has an effect (i.e. it always defaults to true).
  • The backwards compatibility flag --jax_host_callback_ad_transforms introduced in December 2021, has been removed.

• Deprecations:

  • Several jax.numpy APIs have been deprecated following NumPy NEP-52:
    • jax.numpy.NINF has been deprecated. Use -jax.numpy.inf instead.
    • jax.numpy.PZERO has been deprecated. Use 0.0 instead.
    • jax.numpy.NZERO has been deprecated. Use -0.0 instead.
    • jax.numpy.issubsctype(x, t) has been deprecated. Use jax.numpy.issubdtype(x.dtype, t).
    • jax.numpy.row_stack has been deprecated. Use jax.numpy.vstack instead.
    • jax.numpy.in1d has been deprecated. Use jax.numpy.isin instead.
    • jax.numpy.trapz has been deprecated. Use jax.scipy.integrate.trapezoid instead.
  • jax.scipy.linalg.tril and jax.scipy.linalg.triu have been deprecated, following SciPy. Use jax.numpy.tril and jax.numpy.triu instead.
  • jax.lax.prod has been removed after being deprecated in JAX v0.4.11. Use the built-in math.prod instead.
  • A number of exports from jax.interpreters.xla related to defining HLO lowering rules for custom JAX primitives have been deprecated. Custom primitives should be defined using the StableHLO lowering utilities in

... (truncated)

Commits

• d477b92 Start jax and jaxlib 0.4.16 release
• fe0f12e [Pallas] Wire up cost estimates in Mosaic params
• ee846fa Merge pull request #17642 from apaszke:fix-no-teardown
• 9fa96f8 Raise SkipTest before running setUp
• 32ee27b [callbacks] Add support for shardable ordered effects.
• 152af70 Update XLA dependency to use revision
• bcc545a Update XLA dependency to use revision
• f863cfb Relax some test tolerances to fix failures on Linux aarch64.
• bf40f75 Update XLA dependency to use revision
• e78d8a3 [Pallas] Add Mosaic lowering rule for fpowi.
• 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]

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.