jax
2b615135 - [host_callback] Fixes and documentation changes for omnistaging

Commit
5 years ago
[host_callback] Fixes and documentation changes for omnistaging As a nice benefit from omnistaging, the id_print are executed in program order even if the result is not ued. We keep the functions returning values for backwards compatibility, but also so the can be used inside expressions and also so that they can tap values from the backward computation. Omnistaging introduces jaxpr with no invars, which broke host_callback.
Author
Parents
Loading