[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.