Bump jax from 0.2.5 to 0.2.18

[dependabot]

Bumps jax from 0.2.5 to 0.2.18.

Changelog

Sourced from jax's changelog.

jax 0.2.18 (July 21 2021)

• GitHub commits.

• Breaking changes:

  • Support for Python 3.6 has been dropped, per the deprecation policy. Please upgrade to a supported Python version.
  • The minimum jaxlib version is now 0.1.69.
  • The backend argument to {py:func}jax.dlpack.from_dlpack has been removed.

• New features:

  • Added a polar decomposition ({py:func}jax.scipy.linalg.polar).

• Bug fixes:

  • Tightened the checks for lax.argmin and lax.argmax to ensure they are not used with invalid axis value, or with an empty reduction dimension. ({jax-issue}[#7196](https://github.com/google/jax/issues/7196))

jaxlib 0.1.69 (July 9 2021)

• Fix bugs in TFRT CPU backend that results in incorrect results.

jax 0.2.17 (July 9 2021)

• GitHub commits.
• Bug fixes:
  • Default to the older "stream_executor" CPU runtime for jaxlib <= 0.1.68 to work around #7229, which caused wrong outputs on CPU due to a concurrency problem.
• New features:
  • New SciPy function {py:func}jax.scipy.special.sph_harm.
  • Reverse-mode autodiff functions ({func}jax.grad, {func}jax.value_and_grad, {func}jax.vjp, and {func}jax.linear_transpose) support a parameter that indicates which named axes should be summed over in the backward pass if they were broadcasted over in the forward pass. This enables use of these APIs in a non-per-example way inside maps (initially only {func}jax.experimental.maps.xmap) ({jax-issue}[#6950](https://github.com/google/jax/issues/6950)).

jax 0.2.16 (June 23 2021)

• GitHub commits.

jax 0.2.15 (June 23 2021)

• GitHub commits.
• New features:
  • #7042 Turned on TFRT CPU backend with significant dispatch performance improvements on CPU.
  • The {func}jax2tf.convert supports inequalities and min/max for booleans ({jax-issue}[#6956](https://github.com/google/jax/issues/6956)).

... (truncated)

Commits

• fba32c3 Merge pull request #7349 from colemanliyah:cache_size_change
• 86985fd add parameter to initalize_cache function and added init file
• 93a20ee Merge pull request #7346 from skye:version
• a7916f1 Bump jax version and CHANGELOG to 0.2.18
• 8662c5f Merge pull request #7337 from colemanliyah:xmap_and_pjit
• a830dc9 Merge pull request #7343 from LenaMartens:changelist/385967547
• 24c9a93 Add shape and dtype of leaked tracer to UnexpectedTracerError.
• 65c6a6a xmap and pjit integration
• b884cb4 Merge pull request #7338 from hawkinsp:qdwh
• 0dfd76a Remove additional info return value from jax.scipy.linalg.polar().
• 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 #224.