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

[dependabot[bot]]

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

Release notes

Sourced from jax[cpu]'s releases.

JAX release v0.4.13

NOTE: This is the last JAX release that will include Python 3.8 support

• Changes

  • jax.jit now allows None to be passed to in_shardings and out_shardings. The semantics are as follows:
    • For in_shardings, JAX will mark is as replicated but this behavior can change in the future.
    • For out_shardings, we will rely on the XLA GSPMD partitioner to determine the output shardings.
  • jax.experimental.pjit.pjit also allows None to be passed to in_shardings and out_shardings. The semantics are as follows:
    • If the mesh context manager is not provided, JAX has the freedom to choose whatever sharding it wants.
      • For in_shardings, JAX will mark is as replicated but this behavior can change in the future.
      • For out_shardings, we will rely on the XLA GSPMD partitioner to determine the output shardings.
    • If the mesh context manager is provided, None will imply that the value will be replicated on all devices of the mesh.
  • Executable.cost_analysis() works on Cloud TPU
  • Added a warning if a non-allowlisted jaxlib plugin is in use.
  • Added jax.tree_util.tree_leaves_with_path.

• Bug fixes

  • Fixed incorrect wheel name in CUDA 12 releases (#16362); the correct wheel is named cudnn89 instead of cudnn88.

• Deprecations

  • The native_serialization_strict_checks parameter to {func}jax.experimental.jax2tf.convert is deprecated in favor of the new native_serializaation_disabled_checks ({jax-issue}[#16347](https://github.com/google/jax/issues/16347)).

Changelog

Sourced from jax[cpu]'s changelog.

jax 0.4.13 (June 22, 2023)

• Changes

  • jax.jit now allows None to be passed to in_shardings and out_shardings. The semantics are as follows:
    • For in_shardings, JAX will mark is as replicated but this behavior can change in the future.
    • For out_shardings, we will rely on the XLA GSPMD partitioner to determine the output shardings.
  • jax.experimental.pjit.pjit also allows None to be passed to in_shardings and out_shardings. The semantics are as follows:
    • If the mesh context manager is not provided, JAX has the freedom to choose whatever sharding it wants.
      • For in_shardings, JAX will mark is as replicated but this behavior can change in the future.
      • For out_shardings, we will rely on the XLA GSPMD partitioner to determine the output shardings.
    • If the mesh context manager is provided, None will imply that the value will be replicated on all devices of the mesh.
  • Executable.cost_analysis() works on Cloud TPU
  • Added a warning if a non-allowlisted jaxlib plugin is in use.
  • Added jax.tree_util.tree_leaves_with_path.

• Bug fixes

  • Fixed incorrect wheel name in CUDA 12 releases (#16362); the correct wheel is named cudnn89 instead of cudnn88.

• Deprecations

  • The native_serialization_strict_checks parameter to {func}jax.experimental.jax2tf.convert is deprecated in favor of the new native_serializaation_disabled_checks ({jax-issue}[#16347](https://github.com/google/jax/issues/16347)).

jaxlib 0.4.13 (June 22, 2023)

• Changes

  • Added Windows CPU-only wheels to the jaxlib Pypi release.

• Bug fixes

  • __cuda_array_interface__ was broken in previous jaxlib versions and is now fixed ({jax-issue}16440).
  • Concurrent CUDA kernel tracing is now enabled by default on NVIDIA GPUs.

jax 0.4.12 (June 8, 2023)

• Changes

  • Added {class}scipy.spatial.transform.Rotation and {class}scipy.spatial.transform.Slerp

• Deprecations

  • jax.abstract_arrays and its contents are now deprecated. See related functionality in :mod:jax.core.

... (truncated)

Commits

• c3e2427 Merge pull request #16527 from skye:version
• 487b640 Jax 0.4.13 release.
• 10424c5 Update JAX's XlaExecutable.cost_analysis and related plumbing so it works on ...
• 9f4080a Silence pytype errors under an upcoming pytype change.
• e123d1e Merge pull request #16508 from hawkinsp:metal
• 677b0d9 Ignore JAX_USE_PJRT_C_API_ON_TPU=false user warning raised.
• 85a84fd Add a link to the Apple Metal plugin to the JAX README.
• b3527f3 Zlib compress kernel proto.
• f238667 Make JAX-Triton calls serializable.
• c5a47d1 [jax2tf] Refactor the backwards compatibility tests.
• 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[bot]]

Superseded by #830.