Bump jax from 0.2.8 to 0.3.14

[dependabot[bot]]

Bumps jax from 0.2.8 to 0.3.14.

Release notes

Sourced from jax's releases.

JAX release v0.3.14

No release notes provided.

Jaxlib release v0.3.14

No release notes provided.

JAX release v0.3.13

No release notes provided.

Jax release v0.3.12

• Changes
  • Fixes google/jax#10717

Jax release v0.3.11

• Changes
  • {func}jax.lax.eigh now accepts an optional sort_eigenvalues argument that allows users to opt out of eigenvalue sorting on TPU.
• Deprecations
  • Non-array arguments to functions in {mod}jax.lax.linalg are now marked keyword-only. As a backward-compatibility step passing keyword-only arguments positionally yields a warning, but in a future JAX release passing keyword-only arguments positionally will fail. However, most users should prefer to use {mod}jax.numpy.linalg instead.
  • {func}jax.scipy.linalg.polar_unitary, which was a JAX extension to the scipy API, is deprecated. Use {func}jax.scipy.linalg.polar instead.

Jax release v0.3.10

No release notes provided.

Jaxlib release v0.3.10

No release notes provided.

Jax release 0.3.9

• Changes
  • Added support for fully asynchronous checkpointing for GlobalDeviceArray.

JAX release v0.3.8

• GitHub commits.
• Changes
  • {func}jax.numpy.linalg.svd on TPUs uses a qdwh-svd solver.
  • {func}jax.numpy.linalg.cond on TPUs now accepts complex input.
  • {func}jax.numpy.linalg.pinv on TPUs now accepts complex input.
  • {func}jax.numpy.linalg.matrix_rank on TPUs now accepts complex input.
  • {func}jax.scipy.cluster.vq.vq has been added.
  • jax.experimental.maps.mesh has been deleted. Please use jax.experimental.maps.Mesh. Please see https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.maps.Mesh.html#jax.experimental.maps.Mesh for more information.
  • {func}jax.scipy.linalg.qr now returns a length-1 tuple rather than the raw array when mode='r', in order to match the behavior of scipy.linalg.qr ({jax-issue}[#10452](https://github.com/google/jax/issues/10452))
  • {func}jax.numpy.take_along_axis now takes an optional mode parameter that specifies the behavior of out-of-bounds indexing. By default, invalid values (e.g., NaN) will be returned for out-of-bounds indices. In previous versions of JAX, invalid indices were clamped into range. The previous behavior can be restored by passing mode="clip".
  • {func}jax.numpy.take now defaults to mode="fill", which returns invalid values (e.g., NaN) for out-of-bounds indices.

... (truncated)

Changelog

Sourced from jax's changelog.

jax 0.3.14 (June 27, 2022)

• GitHub commits.
• Breaking changes
  • {func}jax.experimental.compilation_cache.initialize_cache does not support max_cache_size_ bytes anymore and will not get that as an input.
  • JAX_PLATFORMS now raises an exception when platform initialization fails.
• Changes
  • Fixed compatibility problems with NumPy 1.23.
  • {func}jax.numpy.linalg.slogdet now accepts an optional method argument that allows selection between an LU-decomposition based implementation and an implementation based on QR decomposition.
  • {func}jax.numpy.linalg.qr now supports mode="raw".
  • pickle, copy.copy, and copy.deepcopy now have more complete support when used on jax arrays ({jax-issue}[#10659](https://github.com/google/jax/issues/10659)). In particular:
    • pickle and deepcopy previously returned np.ndarray objects when used on a DeviceArray; now DeviceArray objects are returned. For deepcopy, the copied array is on the same device as the original. For pickle the deserialized array will be on the default device.
    • Within function transformations (i.e. traced code), deepcopy and copy previously were no-ops. Now they use the same mechanism as DeviceArray.copy().
    • Calling pickle on a traced array now results in an explicit ConcretizationTypeError.
  • The implementation of singular value decomposition (SVD) and symmetric/Hermitian eigendecomposition should be significantly faster on TPU, especially for matrices above 1000x1000 or so. Both now use a spectral divide-and-conquer algorithm for eigendecomposition (QDWH-eig).
  • {func}jax.numpy.ldexp no longer silently promotes all inputs to float64, instead it promotes to float32 for integer inputs of size int32 or smaller ({jax-issue}[#10921](https://github.com/google/jax/issues/10921)).
  • Add a create_perfetto_link option to {func}jax.profiler.start_trace and {func}jax.profiler.start_trace. When used, the profiler will generate a link to the Perfetto UI to view the trace.
  • Changed the semantics of {func}jax.profiler.start_server(...) to store the keepalive globally, rather than requiring the user to keep a reference to it.
  • Added {func}jax.random.generalized_normal.
  • Added {func}jax.random.ball.
  • Added {func}jax.default_device.
  • Added a python -m jax.collect_profile script to manually capture program traces as an alternative to the Tensorboard UI.
  • Added a jax.named_scope context manager that adds profiler metadata to Python programs (similar to jax.named_call).
  • In scatter-update operations (i.e. :attr:jax.numpy.ndarray.at), unsafe implicit dtype casts are deprecated, and now result in a FutureWarning. In a future release, this will become an error. An example of an unsafe implicit cast is jnp.zeros(4, dtype=int).at[0].set(1.5), in which 1.5 previously was silently truncated to 1.
  • {func}jax.experimental.compilation_cache.initialize_cache now supports gcs bucket path as input.
  • Added {func}jax.scipy.stats.gennorm.

... (truncated)

Commits

• 5b576cb Revert: Drop flatbuffers as a Python dependency of JAX.
• 5b865ed Add the assume_metadata option to avoid waiting on ts.open which is very ...
• a2f1aee Merge pull request #11248 from jakevdp:remove-jaxtestcase
• 887abbc jax.test_util: remove deprecated test classes.
• 997beb3 Merge pull request #11273 from hawkinsp:release
• 1e29b7b Update CHANGELOG.md and setup.py for 0.3.14 release.
• 0260360 Merge pull request #11244 from hawkinsp:xla
• f4ddd3e Update XLA.
• efefeac Drop flatbuffers as a Python dependency of JAX.
• 93f5113 Merge pull request #11250 from LenaMartens:changelist/456788427
• 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[bot]]

Superseded by #323.