chore(deps-dev): update jaxlib requirement from <=0.1.65 to <0.3.21

[dependabot[bot]]

Updates the requirements on jaxlib to permit the latest version.

Release notes

Sourced from jaxlib's releases.

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).

Changelog

Sourced from jaxlib's changelog.

jax 0.3.20 (Sep 28, 2022)

• Bug fixes:
  • Adds missing .pyi files that were missing from the previous release ({jax-issue}[#12536](https://github.com/google/jax/issues/12536)).
  • Fixes an incompatibility between jax 0.3.19 and the libtpu version it pinned ({jax-issue}[#12550](https://github.com/google/jax/issues/12550)). Requires jaxlib 0.3.20.
  • Fix incorrect pip url in setup.py comment ({jax-issue}[#12528](https://github.com/google/jax/issues/12528)).

jaxlib 0.3.20 (Sep 28, 2022)

• Bug fixes
  • Fixes support for limiting the visible CUDA devices via jax_cuda_visible_devices in distributed jobs. This functionality is needed for the JAX/SLURM integration on GPU ({jax-issue}[#12533](https://github.com/google/jax/issues/12533)).

jax 0.3.19 (Sep 27, 2022)

• Fixes required jaxlib version.

jax 0.3.18 (Sep 26, 2022)

• 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 overview and the API docs for {mod}jax.stages.
  • Introduced {class}jax.Array, intended to be used for both isinstance checks and type annotations for array types in JAX. Notice that this included some subtle changes to how isinstance works for {class}jax.numpy.ndarray for jax-internal objects, as {class}jax.numpy.ndarray is now a simple alias of {class}jax.Array.
• Breaking changes
  • jax._src is no longer imported into the from the public jax namespace. This may break users that were using JAX internals.
  • jax.soft_pmap has been deleted. Please use pjit or xmap instead. jax.soft_pmap is undocumented. If it were documented, a deprecation period would have been provided.

jax 0.3.17 (Aug 31, 2022)

• GitHub commits.
• Bugs
  • Fix corner case issue in gradient of lax.pow with an exponent of zero ({jax-issue}12041)
• Breaking changes
  • {func}jax.checkpoint, also known as {func}jax.remat, no longer supports the concrete option, following the previous version's deprecation; see JEP 11830.
• Changes
  • Added {func}jax.pure_callback that enables calling back to pure Python functions from compiled functions (e.g. functions decorated with jax.jit or jax.pmap).
• Deprecations:
  • The deprecated DeviceArray.tile() method has been removed. Use {func}jax.numpy.tile ({jax-issue}[#11944](https://github.com/google/jax/issues/11944)).
  • DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead.

jax 0.3.16

... (truncated)

Commits

• 0282b4b Merge pull request #12538 from jakevdp:bundle-pyi
• aafc70d Merge pull request #12556 from hawkinsp:rocm
• 5fe7a54 Merge pull request #12555 from hawkinsp:release
• 39eabe8 Merge pull request #12552 from hawkinsp:nccl
• f7bafb3 Disable multiprocess_gpu_test that fails on ROCm.
• 8d86436 jax/jaxlib 0.3.20 release candidate.
• eabb91e Fix test failure in GPU CI if NCCL_DEBUG is enabled.
• 96abd9a Merge pull request #12540 from sharadmv:cond-lowering-fix
• 96a85bd Make addressable_shards a property like local_shards
• 9489068 Merge pull request #12546 from mattjj:issue12542
• 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)

[kjappelbaum]

will give it a manual look.

[dependabot[bot]]

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.