Fix corner case issue in gradient of lax.pow with an exponent of zero (#12041)
Breaking changes
jax.checkpoint, also known as jax.remat, no longer supports the concrete option, following the previous version's deprecation; see JEP 11830.
Changes
Added jax.pure_callback that enables calling back to pure Python functions from compiled functions (e.g. functions decorated with jax.jit or jax.pmap).
Deprecations:
The deprecated DeviceArray.tile() method has been removed. Use jax.numpy.tile (#11944).
DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead.
Fix corner case issue in gradient of lax.pow with an exponent of zero
({jax-issue}12041)
Breaking changes
{func}jax.checkpoint, also known as {func}jax.remat, no longer supports
the concrete option, following the previous version's deprecation; see
JEP 11830.
Changes
Added {func}jax.pure_callback that enables calling back to pure Python functions from compiled functions (e.g. functions decorated with jax.jit or jax.pmap).
Deprecations:
The deprecated DeviceArray.tile() method has been removed. Use {func}jax.numpy.tile
({jax-issue}[#11944](https://github.com/google/jax/issues/11944)).
DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead.
{func}jax.mask {func}jax.shapecheck APIs have been removed.
See {jax-issue}[#11557](https://github.com/google/jax/issues/11557).
{mod}jax.experimental.loops has been removed. See {jax-issue}[#10278](https://github.com/google/jax/issues/10278)
for an alternative API.
{func}jax.tree_util.tree_multimap has been removed. It has been deprecated since
JAX release 0.3.5, and {func}jax.tree_util.tree_map is a direct replacement.
Removed jax.experimental.stax; it has long been a deprecated alias of
{mod}jax.example_libraries.stax.
Removed jax.experimental.optimizers; it has long been a deprecated alias of
{mod}jax.example_libraries.optimizers.
{func}jax.checkpoint, also known as {func}jax.remat, has a new
implementation switched on by default, meaning the old implementation is
deprecated; see JEP 11830.
JaxTestCase and JaxTestLoader have been removed from jax.test_util. These
classes have been deprecated since v0.3.1 ({jax-issue}[#11248](https://github.com/google/jax/issues/11248)).
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[bot]
commented
2 years ago
Bumps jax from 0.2.8 to 0.3.17.
Release notes
Sourced from jax's releases.
... (truncated)
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
c1192f3Merge pull request #12173 from froystig:random-unwrap0d3630baddkey_datatojax.randomfor key array unwrappingda24b99Some minor changes to make_array_from_callback to use the device_indices_map ...c26c7fdMerge pull request #12167 from froystig:key-dtype39c16c83Rollback of upgrade logistic (sigmoid) function into a lax primitive.52e5d74[Rollback] Make the jaxpr_replicas function explicitly for pmap primitive.cdd16e3Fix typo injax.debug.print()documentation.dc03a33avoid unnecessary lifting to aval injax.random.poisson077bfacadddtypeproperty to key arrays73bf0aaaccess rules through a hidden attribute of opaque dtypeDependabot 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)