Update jax[cpu] requirement from <0.4.20 to <0.4.27

[dependabot]

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

Changelog

Sourced from jax[cpu]'s changelog.

jax 0.4.26 (April 3, 2024)

• New Functionality

  • Added {func}jax.numpy.trapezoid, following the addition of this function in NumPy 2.0.

• Changes

  • Complex-valued {func}jax.numpy.geomspace now chooses the logarithmic spiral branch consistent with that of NumPy 2.0.
  • The behavior of lax.rng_bit_generator, and in turn the 'rbg' and 'unsafe_rbg' PRNG implementations, under jax.vmap has changed so that mapping over keys results in random generation only from the first key in the batch.
  • Docs now use jax.random.key for construction of PRNG key arrays rather than jax.random.PRNGKey.

• Deprecations & Removals

  • {func}jax.tree_map is deprecated; use jax.tree.map instead, or for backward compatibility with older JAX versions, use {func}jax.tree_util.tree_map.
  • {func}jax.clear_backends is deprecated as it does not necessarily do what its name suggests and can lead to unexpected consequences, e.g., it will not destroy existing backends and release corresponding owned resources. Use {func}jax.clear_caches if you only want to clean up compilation caches. For backward compatibility or you really need to switch/reinitialize the default backend, use {func}jax.extend.backend.clear_backends.
  • The jax.experimental.maps module and jax.experimental.maps.xmap are deprecated. Use jax.experimental.shard_map or jax.vmap with the spmd_axis_name argument for expressing SPMD device-parallel computations.
  • The jax.experimental.host_callback module is deprecated. Use instead the new JAX external callbacks. Added JAX_HOST_CALLBACK_LEGACY flag to assist in the transition to the new callbacks. See {jax-issue}[#20385](https://github.com/google/jax/issues/20385) for a discussion.
  • Passing arguments to {func}jax.numpy.array_equal and {func}jax.numpy.array_equiv that cannot be converted to a JAX array now results in an exception.
  • The deprecated flag jax_parallel_functions_output_gda has been removed. This flag was long deprecated and did nothing; its use was a no-op.
  • The previously-deprecated imports jax.interpreters.ad.config and jax.interpreters.ad.source_info_util have now been removed. Use jax.config and jax.extend.source_info_util instead.
  • JAX export does not support older serialization versions anymore. Version 9 has been supported since October 27th, 2023 and has become the default since February 1, 2024. See a description of the versions. This change could break clients that set a specific JAX serialization version lower than 9.

jaxlib 0.4.26 (April 3, 2024)

• Changes

... (truncated)

Commits

• fed7efd Merge pull request #20571 from hawkinsp:release
• 6149326 Prepare for 0.4.26 release.
• 9bb3f79 Avoid unnecessary fori_loop when calculating the block indices.
• 85cb169 Merge pull request #20562 from olupton:bind-to-all
• 57ee6b7 Merge pull request #20560 from pearu:pearu/log1p-fixes
• 1db4af1 Use scipy 1.13.0 instead of scipy 1.13.0rc1 in CI.
• 9a7fb89mpmath/mpmath#774
• e2f4774 Fix tests that fail if enable_checks is true under NumPy 2.0.0rc1.
• 2dd1b3d jax.distributed.initialize: specify bind address.
• d89f0d6 Merge pull request #20534 from gnecula:callback_64bit
• 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)