Bump jax from 0.2.5 to 0.2.26

[dependabot]

Bumps jax from 0.2.5 to 0.2.26.

Release notes

Sourced from jax's releases.

JAX release v0.2.26

• Bug fixes:

• Out-of-bounds indices to jax.ops.segment_sum will now be handled with FILL_OR_DROP semantics, as documented. This primarily afects the reverse-mode derivative, where gradients corresponding to out-of-bounds indices will now be returned as 0. (#8634).

• jax2tf will force the converted code to use XLA for the code fragments under jax.jit, e.g., most jax.numpy functions (#7839).

JAX release v0.2.25

• New features:

  • (Experimental) jax.distributed.initialize exposes multi-host GPU backend.
  • jax.random.permutation supports new independent keyword argument ({jax-issue}[#8430](https://github.com/google/jax/issues/8430))

• Breaking changes

  • Moved jax.experimental.stax to jax.example_libraries.stax
  • Moved jax.experimental.optimizers to jax.example_libraries.optimizers

• New features:

  • Added jax.lax.linalg.qdwh.

Jax release v0.2.24

• New features:
  • jax.random.choice and jax.random.permutation now support multidimensional arrays and an optional axis argument (#8158)
• Breaking changes:
  • jax.numpy.take and jax.numpy.take_along_axis now require array-like inputs (see #7737)

Jax release v0.2.21

• New features:

  • Added jax.numpy.insert implementation (#7936 ).

• Breaking Changes

  • jax.api has been removed. Functions that were available as jax.api.* were aliases for functions in jax.*; please use the functions in jax.* instead.
  • jax.partial, jax.lax.partial, and jax.util.partial were accidental exports that have now been removed. Use functools.partial from the Python standard library instead.
  • Boolean scalar indices now raise a TypeError; previously this silently returned wrong results (#7925 ).
  • Many more jax.numpy functions now require array-like inputs, and will error if passed a list (#7747 #7802 #7907 ). See #7737 for a discussion of the rationale behind this change.
  • When inside a transformation such as jax.jit, jax.numpy.array always stages the array it produces into the traced computation. Previously jax.numpy.array would sometimes produce a on-device array, even under a jax.jit decorator. This change may break code that used JAX arrays to perform shape or index computations that must be known statically; the workaround is to perform such computations using classic NumPy arrays instead.
  • jnp.ndarray is now a true base-class for JAX arrays. In particular, this

... (truncated)

Changelog

Sourced from jax's changelog.

jax 0.2.26 (Dec 8, 2021)

• GitHub commits.

• Bug fixes:

  • Out-of-bounds indices to jax.ops.segment_sum will now be handled with FILL_OR_DROP semantics, as documented. This primarily afects the reverse-mode derivative, where gradients corresponding to out-of-bounds indices will now be returned as 0. (#8634).
  • jax2tf will force the converted code to use XLA for the code fragments under jax.jit, e.g., most jax.numpy functions ({jax-issue}[#7839](https://github.com/google/jax/issues/7839)).

jaxlib 0.1.74 (Nov 17, 2021)

• Enabled peer-to-peer copies between GPUs. Previously, GPU copies were bounced via the host, which is usually slower.
• Added experimental MLIR Python bindings for use by JAX.

jax 0.2.25 (Nov 10, 2021)

• GitHub commits.

• New features:

  • (Experimental) jax.distributed.initialize exposes multi-host GPU backend.
  • jax.random.permutation supports new independent keyword argument ({jax-issue}[#8430](https://github.com/google/jax/issues/8430))

• Breaking changes

  • Moved jax.experimental.stax to jax.example_libraries.stax
  • Moved jax.experimental.optimizers to jax.example_libraries.optimizers

• New features:

  • Added jax.lax.linalg.qdwh.

jax 0.2.24 (Oct 19, 2021)

• GitHub commits.

• New features:

  • jax.random.choice and jax.random.permutation now support multidimensional arrays and an optional axis argument ({jax-issue}[#8158](https://github.com/google/jax/issues/8158))

• Breaking changes:

  • jax.numpy.take and jax.numpy.take_along_axis now require array-like inputs (see {jax-issue}[#7737](https://github.com/google/jax/issues/7737))

jaxlib 0.1.73 (Oct 18, 2021)

• Multiple cuDNN versions are now supported for jaxlib GPU cuda11 wheels.

  • cuDNN 8.2 or newer. We recommend using the cuDNN 8.2 wheel if your cuDNN installation is new enough, since it supports additional functionality.
  • cuDNN 8.0.5 or newer.

• Breaking changes:

... (truncated)

Commits

• 1b5630e Update jaxlib version number to 0.1.76
• a081f00 Merge pull request #8850 from froystig:aot-transform
• 3264045 Restructures jax2tf converters eval directory.
• 6a821d3 Merge pull request #8843 from jakevdp:lax-dtype-errs
• b4db8ab Merge pull request #8849 from jakevdp:x64-linalg
• b980acf detect and err on transformation of AOT-compiled function calls
• 2143f3d Ensure semantically equivalent PartitionSpecs are considered equivalent when
• f8e18e9 [x64] minor weak_type changes to linalg.py
• 4dc3137 Before this change, the local_to_global in the path of dispatching the pjitte...
• 5779267 Merge pull request #8847 from hawkinsp:tol
• 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 #276.