Bump jax from 0.2.8 to 0.3.12

[dependabot[bot]]

Bumps jax from 0.2.8 to 0.3.12.

Release notes

Sourced from jax's releases.

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.
  • Scatter operations, such as x.at[...].set(...), now have "drop" semantics. This has no effect on the scatter operation itself, but it means that when differentiated the gradient of a scatter will yield zero cotangents for out-of-bounds indices. Previously out-of-bounds indices were clamped into range for the gradient, which was not mathematically correct.
  • {func}jax.numpy.take_along_axis now raises a TypeError if its indices are not of an integer type, matching the behavior of {func}numpy.take_along_axis. Previously non-integer indices were silently cast to integers.
  • {func}jax.numpy.ravel_multi_index now raises a TypeError if its dims argument is not of an integer type, matching the behavior of {func}numpy.ravel_multi_index. Previously non-integer dims was silently cast to integers.
  • {func}jax.numpy.split now raises a TypeError if its axis argument is not of an integer type, matching the behavior of {func}numpy.split. Previously non-integer axis was silently cast to integers.
  • {func}jax.numpy.indices now raises a TypeError if its dimensions are not of an integer type, matching the behavior of {func}numpy.indices. Previously non-integer dimensions were silently cast to integers.
  • {func}jax.numpy.diag now raises a TypeError if its k argument is not of an integer type, matching the behavior of {func}numpy.diag. Previously non-integer k was silently cast to integers.
  • Added {func}jax.random.orthogonal.
• Deprecations

... (truncated)

Changelog

Sourced from jax's changelog.

jax 0.3.12 (May 15, 2022)

• GitHub commits.
• Changes
  • Fixes #10717.

jax 0.3.11 (May 15, 2022)

• GitHub commits.
• 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 0.3.10 (May 3, 2022)

• GitHub commits.

jaxlib 0.3.10 (May 3, 2022)

• GitHub commits.
• Changes
  • TF commit fixes an issue in the MHLO canonicalizer that caused constant folding to take a long time or crash for certain programs.

jax 0.3.9 (May 2, 2022)

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

jax 0.3.8 (April 29 2022)

• 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

... (truncated)

Commits

• f572646 Merge pull request #10718 from hawkinsp:jaxfix
• 337ec47 Fix jax 0.3.11 GPU breakge when used with jaxlib 0.3.10.
• 1381afc Update version after jax release
• 42421c8 Merge pull request #10688 from sharadmv:color-by-default
• b8a523f Enable colors when we are using a terminal or IPython
• f26133c Merge pull request #10584 from mattjj:remove-jnp-array-handling-raw-buffers
• 86899ee Merge pull request #10711 from mattjj:closed-call
• 05dda56 add core.closed_call_p
• c0d6a04 remove jnp.array case for handling buffers w/ aval=None
• ba0a2b3 Add profile_version to CompileOptions
• 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 #301.