chore(deps-dev): update jax requirement from <=0.2.12 to <0.3.26

[dependabot]

Updates the requirements on jax to permit the latest version.

Changelog

Sourced from jax's changelog.

jax 0.3.25 (Nov 15, 2022)

• Changes
  • {func}jax.numpy.linalg.pinv now supports the hermitian option.
  • {func}jax.scipy.linalg.hessenberg is now supported on CPU only. Requires jaxlib > 0.3.24.
  • New functions {func}jax.lax.linalg.hessenberg, {func}jax.lax.linalg.tridiagonal, and {func}jax.lax.linalg.householder_product were added. Householder reduction is currently CPU-only and tridiagonal reductions are supported on CPU and GPU only.
  • The gradients of svd and jax.numpy.linalg.pinv are now computed more economically for non-square matrices.
• Breaking Changes
  • Deleted the jax_experimental_name_stack config option.
  • Convert a string axis_names arguments to the {class}jax.experimental.maps.Mesh constructor into a singleton tuple instead of unpacking the string into a sequence of character axis names.

jaxlib 0.3.25 (Nov 15, 2022)

• Changes
  • Added support for tridiagonal reductions on CPU and GPU.
  • Added support for upper Hessenberg reductions on CPU.
• Bugs
  • Fixed a bug that meant that frames in tracebacks captured by JAX were incorrectly mapped to source lines under Python 3.10+

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)

... (truncated)

Commits

• a683186 Use the 11/09 libtpu build for jaxlib release since that passes all the tests.
• 7635df8 Remove custom potrf kernels in favor of native XLA cholesky support.
• f36084a Update the values for jaxlib release (again)
• 0da02dd Update the values needed for a jaxlib release
• 465d4ec Merge pull request #13243 from google:yashk2810-patch-19
• 50f7512 Update WORKSPACE for jaxlib release
• e6c4d4a Add docstrings for Sharding classes. Right now I am only documenting `Shard...
• c42bad8 Make MeshPspecSharding an alias for NamedSharding (it was the other way a...
• d1e26d9 Merge pull request #13139 from mattjj:djax-vmap4
• f715ca6 Merge pull request #13236 from skye:build_chat
• 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]

OK, I won't notify you again about this release, but will get in touch when a new version is available. If you'd rather skip all updates until the next major or minor version, let me know by commenting @dependabot ignore this major version or @dependabot ignore this minor version. You can also ignore all major, minor, or patch releases for a dependency by adding an ignore condition with the desired update_types to your config file.

If you change your mind, just re-open this PR and I'll resolve any conflicts on it.