Bump jax from 0.2.5 to 0.2.22

[dependabot]

Bumps jax from 0.2.5 to 0.2.22.

Release notes

Sourced from jax's releases.

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 means that for a standard numpy array x, isinstance(x, jnp.ndarray) will now return False (#7927).

Jax release v0.2.20

No release notes provided.

Changelog

Sourced from jax's changelog.

jax 0.2.22 (Oct 12, 2021)

• GitHub commits.
• Breaking Changes

  • Static arguments to jax.pmap must now be hashable.

    Unhashable static arguments have long been disallowed on jax.jit, but they were still permitted on jax.pmap; jax.pmap compared unhashable static arguments using object identity.

    This behavior is a footgun, since comparing arguments using object identity leads to recompilation each time the object identity changes. Instead, we now ban unhashable arguments: if a user of jax.pmap wants to compare static arguments by object identity, they can define __hash__ and __eq__ methods on their objects that do that, or wrap their objects in an object that has those operations with object identity semantics. Another option is to use functools.partial to encapsulate the unhashable static arguments into the function object.

  • jax.util.partial was an accidental export that has now been removed. Use functools.partial from the Python standard library instead.

• Deprecations
  • The functions jax.ops.index_update, jax.ops.index_add etc. are deprecated and will be removed in a future JAX release. Please use the .at property on JAX arrays instead, e.g., x.at[idx].set(y). For now, these functions produce a DeprecationWarning.
• New features:
  • An optimized C++ code-path improving the dispatch time for pmap is now the default when using jaxlib 0.1.72 or newer. The feature can be disabled using the --experimental_cpp_pmap flag (or JAX_CPP_PMAP environment variable).
  • jax.numpy.unique now supports an optional fill_value argument ({jax-issue}[#8121](https://github.com/google/jax/issues/8121))

jaxlib 0.1.72 (Oct 12, 2021)

• Breaking changes:
  • Support for CUDA 10.2 and CUDA 10.1 has been dropped. Jaxlib now supports CUDA 11.1+.

jax 0.2.21 (Sept 23, 2021)

• GitHub commits.
• 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, and jax.lax.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 ({jax-issue}[#7925](https://github.com/google/jax/issues/7925)).
  • Many more jax.numpy functions now require array-like inputs, and will error

... (truncated)

Commits

• 75062c1 Merge pull request #8188 from skye:version
• 962c496 Update jax version and CHANGELOG for 0.2.22 release
• 66a4a9f Remove 10.2 cuda support
• 7a3bf61 Merge pull request #8187 from skye:version
• 0072c32 Update CHANGELOG and verson numbers for jaxlib 0.1.72 release
• 075d835 Merge pull request #8183 from froystig:named-shape-hash
• 9a182e6 order-independent hash in core.NamedShape
• f34387b work around unhashable named shapes in api.ShapeDtypeStruct
• b30fa13 Merge pull request #8145 from froystig:aot-pjit
• 8f92477 Merge pull request #8168 from froystig:prng-raw
• 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 #250.