Jax

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.2.3...jax-v0.2.4).
* Improvements:
* Add support for `remat` to jax.experimental.host_callback. See {jax-issue}`4608`.
* Deprecations

* Indexing with non-tuple sequences is now deprecated, following a similar deprecation in Numpy.
In a future release, this will result in a TypeError. See {jax-issue}`4564`.

jaxlib 0.1.56 (October 14, 2020)

* [GitHub commits](https://github.com/google/jax/compare/jax-v0.2.2...jax-v0.2.3).
* The reason for another release so soon is we need to temporarily roll back a
new jit fastpath while we look into a performance degradation

* [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)