Support for NumPy 1.18 has been dropped, per the [deprecation policy](https://jax.readthedocs.io/en/latest/ deprecation.html). Please upgrade to a supported NumPy version.
The host_callback primitives have been simplified to drop the special autodiff handling for hcb.id_tap and id_print. From now on, only the primals are tapped. The old behavior can be obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS environment variable, or the --flax_host_callback_ad_transforms flag. Additionally, added documentation for how to implement the old behavior using JAX custom AD APIs ({jax-issue}[#8678](https://github.com/google/jax/issues/8678)).
Sorting now matches the behavior of NumPy for 0.0 and NaN regardless of the bit representation. In particular, 0.0 and -0.0 are now treated as equivalent, where previously -0.0 was treated as less than 0.0. Additionally all NaN representations are now treated as equivalent and sorted to the end of the array. Previously negative NaN values were sorted to the front of the array, and NaN values with different internal bit representations were not treated as equivalent, and were sorted according to those bit patterns ({jax- issue}[#9178](https://github.com/google/jax/issues/9178)).
{func}jax.numpy.unique now treats NaN values in the same way as np.unique in NumPy versions 1.21 and newer: at most one NaN value will appear in the uniquified output ({jax-issue}9184).
Bug fixes:
host_callback now supports ad_checkpoint.checkpoint ({jax-issue}[#8907](https://github.com/google/jax/issues/8907)).
Added a new debugging flag/environment variable JAX_DUMP_IR_TO=/path. If set, JAX dumps the MHLO/HLO IR it generates for each computation to a file under the given path.
Added jax.ensure_compile_time_eval to the public api ({jax-issue}[#7987](https://github.com/google/jax/issues/7987)).
jax2tf now supports a flag jax2tf_associative_scan_reductions to change the lowering for associative reductions, e.g., jnp.cumsum, to behave like JAX on CPU and GPU (to use an associative scan). See the jax2tf README for more details ({jax-issue}[#9189](https://github.com/google/jax/issues/9189)).
JAX release v0.2.26
Bug fixes:
Out-of-bounds indices to jax.ops.segment_sum will now be handled with FILL_OR_DROP semantics, as documented. This primarily afects the reverse-mode derivative, where gradients corresponding to out-of-bounds indices will now be returned as 0. (#8634).
jax2tf will force the converted code to use XLA for the code fragments under jax.jit, e.g., most jax.numpy functions (#7839).
Support for NumPy 1.18 has been dropped, per the
deprecation policy.
Please upgrade to a supported NumPy version.
The host_callback primitives have been simplified to drop the
special autodiff handling for hcb.id_tap and id_print.
From now on, only the primals are tapped. The old behavior can be
obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS
environment variable, or the --flax_host_callback_ad_transforms flag.
Additionally, added documentation for how to implement the old behavior
using JAX custom AD APIs ({jax-issue}[#8678](https://github.com/google/jax/issues/8678)).
Sorting now matches the behavior of NumPy for 0.0 and NaN regardless of the
bit representation. In particular, 0.0 and -0.0 are now treated as equivalent,
where previously -0.0 was treated as less than 0.0. Additionally all NaN
representations are now treated as equivalent and sorted to the end of the array.
Previously negative NaN values were sorted to the front of the array, and NaN
values with different internal bit representations were not treated as equivalent, and
were sorted according to those bit patterns ({jax-issue}[#9178](https://github.com/google/jax/issues/9178)).
{func}jax.numpy.unique now treats NaN values in the same way as np.unique in
NumPy versions 1.21 and newer: at most one NaN value will appear in the uniquified
output ({jax-issue}9184).
Bug fixes:
host_callback now supports ad_checkpoint.checkpoint ({jax-issue}[#8907](https://github.com/google/jax/issues/8907)).
Added a new debugging flag/environment variable JAX_DUMP_IR_TO=/path.
If set, JAX dumps the MHLO/HLO IR it generates for each computation to a
file under the given path.
Added jax.ensure_compile_time_eval to the public api ({jax-issue}[#7987](https://github.com/google/jax/issues/7987)).
jax2tf now supports a flag jax2tf_associative_scan_reductions to change
the lowering for associative reductions, e.g., jnp.cumsum, to behave
like JAX on CPU and GPU (to use an associative scan). See the jax2tf README
for more details ({jax-issue}[#9189](https://github.com/google/jax/issues/9189)).
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.2.27.
Release notes
Sourced from jax's releases.
... (truncated)
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
4c423c3Speed up check_jaxpr().e30b96cMerge pull request #9201 from LenaMartens:changelist/4207945528ea8576Checkify: add way to disable categories of errors.6411f8aMerge pull request #9184 from jakevdp:unique-nanbebe984Merge pull request #9205 from jakevdp:einsum-tupleb92db58Canonicalize parsed partition spec before passing to lower_mesh_computation. ...b509aaeSplit lax_control_flow_test into three separate tests.c9169faMerge pull request #9189 from gnecula:tf_reduce_window77d60cfeinsum: clarify use of precision.7f07f1bMerge pull request #9200 from google:LenaMartens-patch-1Dependabot 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)