Bump jax from 0.2.5 to 0.3.24

[dependabot]

Bumps jax from 0.2.5 to 0.3.24.

Release notes

Sourced from jax's releases.

Jax release v0.3.24

No release notes provided.

Jaxlib release v0.3.24

No release notes provided.

jax-v0.3.23

• Changes
  • Update Colab TPU driver version for new jaxlib release.

jax-v0.3.22

• Changes
  • Add JAX_PLATFORMS=tpu,cpu as default setting in TPU initialization, so JAX will raise an error if TPU cannot be initialized instead of falling back to CPU. Set JAX_PLATFORMS='' to override this behavior and automatically choose an available backend (the original default), or set JAX_PLATFORMS=cpu to always use CPU regardless of if the TPU is available.
• Deprecations
  • Several test utilities deprecated in JAX v0.3.8 are now removed from {mod}jax.test_util.

jaxlib-v0.3.22

No release notes provided.

JAX release v0.3.21

• Changes
  • The persistent compilation cache will now warn instead of raising an exception on error ({jax-issue}[#12582](https://github.com/google/jax/issues/12582)), so program execution can continue if something goes wrong with the cache. Set JAX_RAISE_PERSISTENT_CACHE_ERRORS=true to revert this behavior.

JAX release v0.3.20

Notable changes:

• Adds missing .pyi files that were missing from the previous release (#12536).
• Fixes an incompatibility between jax 0.3.19 and the libtpu version it pinned (#12550). Requires jaxlib 0.3.20.
• Fix incorrect pip url in setup.py comment (#12528).

jaxlib release v0.3.20

Notable changes:

• Fixes support for limiting the visible CUDA devices viajax_cuda_visible_devices in distributed jobs. This functionality is needed for the JAX/SLURM integration on GPU (#12533).

JAX release v0.3.19

Fixes the required jaxlib version

JAX release v0.3.18

• GitHub commits.
• Changes
  • Ahead-of-time lowering and compilation functionality (tracked in {jax-issue}[#7733](https://github.com/google/jax/issues/7733)) is stable and public. See [the

... (truncated)

Changelog

Sourced from jax's changelog.

jax 0.3.24 (Nov 4, 2022)

• Changes
  • JAX should be faster to import. We now import scipy lazily, which accounted for a significant fraction of JAX's import time.
  • Setting the env var JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=$N can be used to limit the number of cache entries written to the persistent cache. By default, computations that take 1 second or more to compile will be cached.
    • Added {func}jax.scipy.stats.mode.
  • The default device order used by pmap on TPU if no order is specified now matches jax.devices() for single-process jobs. Previously the two orderings differed, which could lead to unnecessary copies or out-of-memory errors. Requiring the orderings to agree simplifies matters.
• Breaking Changes
  • {func}jax.numpy.gradient now behaves like most other functions in {mod}jax.numpy, and forbids passing lists or tuples in place of arrays ({jax-issue}[#12958](https://github.com/google/jax/issues/12958))
  • Functions in {mod}jax.numpy.linalg and {mod}jax.numpy.fft now uniformly require inputs to be array-like: i.e. lists and tuples cannot be used in place of arrays. Part of {jax-issue}[#7737](https://github.com/google/jax/issues/7737).
• Deprecations
  • jax.sharding.MeshPspecSharding has been renamed to jax.sharding.NamedSharding. jax.sharding.MeshPspecSharding name will be removed in 3 months.

jaxlib 0.3.24 (Nov 4, 2022)

• Changes
  • Buffer donation now works on CPU. This may break code that marked buffers for donation on CPU but relied on donation not being implemented.

jax 0.3.23 (Oct 12, 2022)

• Changes
  • Update Colab TPU driver version for new jaxlib release.

jax 0.3.22 (Oct 11, 2022)

• Changes
  • Add JAX_PLATFORMS=tpu,cpu as default setting in TPU initialization, so JAX will raise an error if TPU cannot be initialized instead of falling back to CPU. Set JAX_PLATFORMS='' to override this behavior and automatically choose an available backend (the original default), or set JAX_PLATFORMS=cpu to always use CPU regardless of if the TPU is available.
• Deprecations
  • Several test utilities deprecated in JAX v0.3.8 are now removed from {mod}jax.test_util.

jaxlib 0.3.22 (Oct 11, 2022)

jax 0.3.21 (Sep 30, 2022)

• GitHub commits.
• Changes
  • The persistent compilation cache will now warn instead of raising an exception on error ({jax-issue}[#12582](https://github.com/google/jax/issues/12582)), so program execution can continue

... (truncated)

Commits

• 46368e4 [sparse] Update the guard of cusparse SpMM and SpMv algorithms to cusparse ve...
• 974134f Merge pull request #13103 from mattjj:issue13099
• 4033007 improve error when f_vjp gets more than one argument
• 2e384ce Prepare for release of jax and jaxlib 0.3.24
• 2b4735f Merge pull request #13094 from mattjj:fix-random-docs-table
• dba9fc0 Merge pull request #13000 from skye:self_hosted_tpu
• 8c22e34 Add Github Actions workflow that runs on a self-hosted TPU VM runner.
• 478bd3e fix comparison table in random docs
• f4be5ab Merge pull request #12219 from jakevdp:indexing-slice
• 532cd7e Skip the benchmarks properly via state.skip_with_error when enough devices ar...
• 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 #396.