Bump jax from 0.2.7 to 0.2.8

[dependabot[bot]]

Bumps jax from 0.2.7 to 0.2.8.

Changelog

Sourced from jax's changelog.

jax 0.2.8 (January 12 2021)

• GitHub commits <https://github.com/google/jax/compare/jax-v0.2.7...jax-v0.2.8>_.

• New features:

  • Add :func:jax.closure_convert for use with higher-order custom derivative functions. ([#5244](https://github.com/google/jax/issues/5244) <https://github.com/google/jax/pull/5244>_)
  • Add :func:jax.experimental.host_callback.call to call a custom Python function on the host and return a result to the device computation. ([#5243](https://github.com/google/jax/issues/5243) <https://github.com/google/jax/pull/5243>_)

• Bug fixes:

  • jax.numpy.arccosh now returns the same branch as numpy.arccosh for complex inputs ([#5156](https://github.com/google/jax/issues/5156) <https://github.com/google/jax/pull/5156>_)
  • host_callback.id_tap now works for jax.pmap also. There is a optional parameter for id_tap and id_print to request that the device from which the value is tapped be passed as a keyword argument to the tap function ([#5182](https://github.com/google/jax/issues/5182) <https://github.com/google/jax/pull/5182>_).

• Breaking changes:

  • jax.numpy.pad now takes keyword arguments. Positional argument constant_values has been removed. In addition, passing unsupported keyword arguments raises an error.

  • Changes for :func:jax.experimental.host_callback.id_tap ([#5243](https://github.com/google/jax/issues/5243) <https://github.com/google/jax/pull/5243>_):

    • Removed support for kwargs for :func:jax.experimental.host_callback.id_tap. (This support has been deprecated for a few months.)
    • Changed the printing of tuples for :func:jax.experimental.host_callback.id_print to use '(' instead of '['.
    • Changed the :func:jax.experimental.host_callback.id_print in presence of JVP to print a pair of primal and tangent. Previously, there were two separate print operations for the primals and the tangent.
    • host_callback.outfeed_receiver has been removed (it is not necessary, and was deprecated a few months ago).

Commits

• c42fc28 Merge pull request #5379 from hawkinsp:jaxrelease2
• 994db37 Release jax 0.2.8.
• 3ac809e [JAX] Move jax.util to jax._src_util.
• bba98a2 Re-enable test.
• a87978f Enable more TPU tests
• d9cd5e3 Merge pull request #5369 from skye:workspace
• 345267b Update TF version for jaxlib 0.1.58 release
• 2ec9b13 Add workaround for pybind11 bug that means we were issue poor errors in the c...
• 4a36d73 Merge pull request #5363 from hawkinsp:xla
• 7d78d1e Update XLA.
• 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)