Bump jax from 0.4.13 to 0.4.14 in /requirements

[dependabot]

Bumps jax from 0.4.13 to 0.4.14.

Changelog

Sourced from jax'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

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 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 #382.