pyro-ppl / numpyro / 39082222312

93%

push

github

perf(util): avoid per-step host sync from the diagnostics string in the MCMC progress-bar loop (#2317)

* Avoid per-step host sync from the diagnostics string in the MCMC progress-bar loop

The single-chain progress-bar loop in fori_collect formatted the kernel
diagnostics (device scalars) on every iteration, which blocks the host
until the current step has finished and defeats JAX's async dispatch for
the whole run. tqdm only displays the postfix every progress_rate steps,
so compute the string only on those iterations and on the last step.
Samples are unchanged.

* Gate the diagnostics on tqdm's live miniters and test fori_collect directly

* Skip the diagnostics on a disabled tqdm bar and guard progress_rate=0

3 of 3 new or added lines in 1 file covered. (100.0%)

15694 of 16913 relevant lines covered (92.79%)

1.7 hits per line