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).
New features:
New flag for debugging infs, analagous to that for NaNs ([#5224](https://github.com/google/jax/issues/5224) <https://github.com/google/jax/pull/5224>_).
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)
Bumps jaxlib from 0.1.58 to 0.1.59.
Changelog
Sourced from jaxlib's changelog.
... (truncated)
Commits
c5d6281
Merge pull request #5430 from skye:fix0927c29
Disable another test that requires omnistagingf79182c
Merge pull request #5429 from jakevdp:disable-broken-test26262cb
disable test that requires omnistaging670f537
Merge pull request #5425 from hawkinsp:typing1abf383
Remove type overloads for tree_map and tree_multimap.500849d
Merge pull request #5424 from skye:workspace62b864a
Update WORKSPACE and jaxlib version to 0.1.596e19774
Merge pull request #5423 from hawkinsp:typedf13a63f
Merge pull request #5422 from gnecula:tap_no_argsDependabot 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)