Bump jax from 0.4.14 to 0.4.21

[dependabot]

Bumps jax from 0.4.14 to 0.4.21.

Changelog

Sourced from jax's changelog.

jax 0.4.21 (Dec 4 2023)

• New Features

  • Added {obj}jax.nn.squareplus.

• Changes

  • The minimum jaxlib version is now 0.4.19.
  • Released wheels are built now with clang instead of gcc.
  • Enforce that the device backend has not been initialized prior to calling jax.distributed.initialize().
  • Automate arguments to jax.distributed.initialize() in cloud TPU environments.

• Deprecations

  • The previously-deprecated sym_pos argument has been removed from {func}jax.scipy.linalg.solve. Use assume_a='pos' instead.
  • Passing None to {func}jax.array or {func}jax.asarray, either directly or within a list or tuple, is deprecated and now raises a {obj}FutureWarning. It currently is converted to NaN, and in the future will raise a {obj}TypeError.
  • Passing the condition, x, and y parameters to jax.numpy.where by keyword arguments has been deprecated, to match numpy.where.
  • Passing arguments to {func}jax.numpy.array_equal and {func}jax.numpy.array_equiv that cannot be converted to a JAX array is deprecated and now raises a {obj}DeprecationWaning. Currently the functions return False, in the future this will raise an exception.
  • The device() method of JAX arrays deprecated. Depending on the context, it may be replaced with one of the following:
    • {meth}jax.Array.devices returns the set of all devices used by the array.
    • {attr}jax.Array.sharding gives the sharding configuration used by the array.

jaxlib 0.4.21 (Dec 4 2023)

• Changes

  • In preparation for adding distributed CPU support, JAX now treats CPU devices identically to GPU and TPU devices, that is:

    • jax.devices() includes all devices present in a distributed job, even those not local to the current process. jax.local_devices() still only includes devices local to the current process, so if the change to jax.devices() breaks you, you most likely want to use jax.local_devices() instead.
    • CPU devices now receive a globally unique ID number within a distributed job; previously CPU devices would receive a process-local ID number.
    • The process_index of each CPU device will now match any GPU or TPU devices within the same process; previously the process_index of a CPU device was always 0.

  • On NVIDIA GPU, JAX now prefers a Jacobi SVD solver for matrices up to 1024x1024. The Jacobi solver appears faster than the non-Jacobi version.

• Bug fixes

  • Fixed error/hang when an array with non-finite values is passed to a

... (truncated)

Commits

• 5942e15 Prepare for 0.4.21 release
• d91c13e Merge pull request #18795 from gnecula:test_export_grad
• 8a2d4a0 [export] Add and fix a test for exporting higher-order gradients with sharding
• 1d95e79 Disable export_harnesses_multi_platform_test under sanitizers.
• 5e0993c Merge pull request #18794 from olupton:qualname
• 3c0c6b7 Use qualified name if possible.
• a137edc Update XLA dependency to use revision
• bf08411 Update XLA dependency to use revision
• 61e79cd Merge pull request #18786 from gnecula:test_export_effects
• b51b80e Merge pull request #18761 from gnecula:export_sharding
• 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]

Superseded by #31.