Update jax[cpu] requirement from <0.4.13 to <0.4.15

[dependabot]

Updates the requirements on jax[cpu] to permit the latest version.

Changelog

Sourced from jax[cpu]'s changelog.

jax 0.4.14 (July 27, 2023)

• Changes

  • jax.jit takes donate_argnames as an argument. It's semantics are similar to static_argnames. If neither donate_argnums nor donate_argnames is provided, no arguments are donated. If donate_argnums is not provided but donate_argnames is, or vice versa, JAX uses inspect.signature(fun) to find any positional arguments that correspond to donate_argnames (or vice versa). If both donate_argnums and donate_argnames are provided, inspect.signature is not used, and only actual parameters listed in either donate_argnums or donate_argnames will be donated.
  • {func}jax.random.gamma has been re-factored to a more efficient algorithm with more robust endpoint behavior ({jax-issue}[#16779](https://github.com/google/jax/issues/16779)). This means that the sequence of values returned for a given key will change between JAX v0.4.13 and v0.4.14 for gamma and related samplers (including {func}jax.random.ball, {func}jax.random.beta, {func}jax.random.chisquare, {func}jax.random.dirichlet, {func}jax.random.generalized_normal, {func}jax.random.loggamma, {func}jax.random.t).

• Deletions

  • in_axis_resources and out_axis_resources have been deleted from pjit since it has been more than 3 months since their deprecation. Please use in_shardings and out_shardings as the replacement. This is a safe and trivial name replacement. It does not change any of the current pjit semantics and doesn't break any code. You can still pass in PartitionSpecs to in_shardings and out_shardings.

• Deprecations

  • Python 3.8 support has been dropped as per https://jax.readthedocs.io/en/latest/deprecation.html
  • JAX now requires NumPy 1.22 or newer as per https://jax.readthedocs.io/en/latest/deprecation.html
  • Passing optional arguments to {func}jax.numpy.ndarray.at by position is no longer supported, after being deprecated in JAX version 0.4.7. For example, instead of x.at[i].get(True), use x.at[i].get(indices_are_sorted=True)
  • The following jax.Array methods have been removed, after being deprecated in JAX v0.4.5:
    • jax.Array.broadcast: use {func}jax.lax.broadcast instead.
    • jax.Array.broadcast_in_dim: use {func}jax.lax.broadcast_in_dim instead.
    • jax.Array.split: use {func}jax.numpy.split instead.
  • The following APIs have been removed after previous deprecation:
    • jax.ad: use {mod}jax.interpreters.ad.
    • jax.curry: use curry = lambda f: partial(partial, f).
    • jax.partial_eval: use {mod}jax.interpreters.partial_eval.
    • jax.pxla: use {mod}jax.interpreters.pxla.
    • jax.xla: use {mod}jax.interpreters.xla.
    • jax.ShapedArray: use {class}jax.core.ShapedArray.
    • jax.interpreters.pxla.device_put: use {func}jax.device_put.
    • jax.interpreters.pxla.make_sharded_device_array: use {func}jax.make_array_from_single_device_arrays.

... (truncated)

Commits

• 88a60b8 Merge pull request #16870 from skye:version
• e132a0e Slightly downgrade xla version to avoid PJRT C API incompat
• c75e85d Merge pull request #16869 from skye:version
• a480aa8 Work around pytype error.
• 0b24b2b Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.14 release
• 7c0ef86 Merge pull request #16433 from JGameCreation:patch-1
• 76cda0a Update flags to use the ABSL typed flag API.
• f35f226 Merge pull request #16865 from google:xla
• 0cd0253 Update XLA version to fix build failure.
• a03d6e6 Move _tpu_ext.cc to jaxlib/mlir/_mlir_libs and set RPATH correctly
• Additional commits viewable in compare view

You can trigger a rebase of this PR 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 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)

Note
Automatic rebases have been disabled on this pull request as it has been open for over 30 days.