pyro-ppl / numpyro / 32657710378

92%

push

github

Avoid per-step host sync in the SVI progress-bar loop (#2247)

The progress-bar loop in SVI.run called jax.device_get(loss) on every
step, blocking on the device each iteration and defeating async
dispatch. Collect losses as device arrays and transfer each display
batch with a single device_get instead; also hoist jit(body_fn) out of
the loop. 1.9x faster on a dispatch-bound model (CPU); results are
bit-identical.

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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

14633 of 15841 relevant lines covered (92.37%)

1.69 hits per line