Jax

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.2.1...jax-v0.2.2).

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.2.0...jax-v0.2.1).
* Improvements:
* As a benefit of omnistaging, the host_callback functions are executed (in program
order) even if the result of the {py:func}`jax.experimental.host_callback.id_print`/
{py:func}`jax.experimental.host_callback.id_tap` is not used in the computation.

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.1.77...jax-v0.2.0).
* Improvements:
* Omnistaging on by default. See {jax-issue}`3370` and
[omnistaging](https://github.com/google/jax/blob/main/docs/design_notes/omnistaging.md)

* Breaking changes:
* New simplified interface for {py:func}`jax.experimental.host_callback.id_tap` (4101)

jaxlib 0.1.55 (September 8, 2020)

* Update XLA:
* Fix bug in DLPackManagedTensorToBuffer (4196)

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.1.75...jax-v0.1.76).

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.1.74...jax-v0.1.75).
* Bug Fixes:
* make jnp.abs() work for unsigned inputs (3914)
* Improvements:
* "Omnistaging" behavior added behind a flag, disabled by default (3370)