Avoid per-step host sync in the SVI progress-bar loop

[kyo219]

Changes made

• numpyro/infer/svi.py: the progress-bar loop in SVI.run called jax.device_get(loss) on every step, forcing a blocking device→host transfer per iteration — the host cannot dispatch step i+1 until step i has finished computing, so JAX's async dispatch is defeated for the whole run. Losses are now appended as device arrays and each display batch (num_steps // 20 steps, the existing postfix-update granularity) is fetched with a single jax.device_get, which also replaces the device scalars with host values as before so memory behavior is unchanged. jit(body_fn) is also hoisted out of the loop instead of being re-wrapped on every iteration.

Results are unchanged: the losses returned by the progress-bar path are bit-identical to master (verified locally on a 1234-step run), and the postfix string is computed from the same values as before.

Benchmarks

Apple M2 (24 GB), CPU, jax 0.10.2, Python 3.11. Each variant runs in a fresh process; time is the best of 3 repeats after a warmup run of svi.run(..., progress_bar=True).

scenario	master	this PR	progress_bar=False (reference)
dispatch-bound: linear regression, N=100, D=3, 20,000 steps	1.90 s	1.00 s (1.9×)	0.28 s
compute-bound: linear regression, N=200,000, D=100, 2,000 steps	12.96 s	12.63 s	11.99 s

The win comes from removing the per-step synchronization, so it scales with the ratio of dispatch overhead to per-step compute; on accelerators, where a device round-trip is more expensive than on CPU, the effect should be larger.

Links to related issues/PRs

None.

Tests

• New test_run_progress_bar_matches_scan_path in test/infer/test_svi.py — locks in that the progress-bar path returns the same losses and params as the lax.scan path, using a step count (123) that is not a multiple of the display batch so the tail losses collected outside a batch boundary are exercised. Tolerance is rtol=1e-5 because the two paths run differently compiled programs and already differ by a few ulps on master.
• pytest test/infer/test_svi.py — 66 passed, 4 skipped
• ruff check / ruff format --check clean; ty check reports only the diagnostic already present on master.

Dependencies

None.

🤖 Generated with Claude Code

[Qazalbash]

Thanks @kyo219

[github-actions]

Benchmark report

this PR svi-progress-bar-host-sync at 305f8c3f vs baseline master at 999d8d1f

  run time:     unchanged across 32 benchmarks
+ compile time: 1 faster

Significant changes (1)

                                ──────── run time ───────     ────── compile time ──────
  benchmark                     baseline   this PR      Δ     baseline   this PR       Δ
────────────────────────────────────────────────────────────────────────────────────────
+ predictive_forward_sampling   702.5 ms  703.7 ms  +0.2%     208.6 ms  155.7 ms  -25.4%

_{Red is slower, green is faster; a row is coloured by the worse of its two columns. A delta in parentheses cleared the threshold on a measurement below the resolution floor, so it is shown without being called a change. † marks a benchmark that could not be compared — see below.}

Full results

distributions

                                 ──────── run time ────────     ────── compile time ─────
  benchmark                      baseline  this PR        Δ     baseline   this PR      Δ
─────────────────────────────────────────────────────────────────────────────────────────
  biject_to_constraints            3.8 ms   3.9 ms    +3.3%     354.7 ms  368.8 ms  +4.0%
  categorical_log_prob             2.2 ms   2.2 ms    -0.5%      66.4 ms   69.3 ms  +4.3%
  dirichlet_log_prob               616 µs   616 µs    -0.1%     382.7 ms  413.9 ms  +8.2%
  dirichlet_sample                56.1 ms  56.2 ms    +0.1%     824.8 ms  806.5 ms  -2.2%
  gamma_log_prob                   2.0 ms   2.1 ms    +3.8%       1.98 s    2.03 s  +2.9%
  gamma_sample                    26.4 ms  26.3 ms    -0.4%     766.2 ms  783.2 ms  +2.2%
  lkj_cholesky_sample              5.6 ms   5.7 ms    +0.5%       1.17 s    1.11 s  -4.8%
  mixture_same_family_log_prob     2.1 ms   2.1 ms    -1.8%     104.0 ms  100.4 ms  -3.4%
  multivariate_normal_log_prob     264 µs   248 µs  (-6.1%)     147.1 ms  156.4 ms  +6.3%
  normal_log_prob                  546 µs   547 µs    +0.2%      53.4 ms   54.4 ms  +1.9%
  normal_sample                   20.9 ms  21.2 ms    +1.4%     191.3 ms  193.2 ms  +1.0%
  stick_breaking_transform         6.2 ms   6.4 ms    +4.5%     207.1 ms  197.6 ms  -4.6%
  student_t_log_prob               3.0 ms   3.0 ms    +0.6%      74.4 ms   74.7 ms  +0.3%
  truncated_normal_log_prob        575 µs   553 µs    -3.7%      51.6 ms   51.5 ms  -0.1%

handlers

                                  ──────── run time ───────     ────── compile time ──────
  benchmark                       baseline   this PR      Δ     baseline   this PR       Δ
──────────────────────────────────────────────────────────────────────────────────────────
  initialize_model_hierarchical    36.7 ms   36.8 ms  +0.2%       3.54 s    3.86 s   +9.0%
  log_density_hierarchical          2.9 ms    3.0 ms  +1.6%       1.13 s    1.14 s   +0.4%
  nested_handler_stack              1.2 ms    1.2 ms  +0.7%       762 µs    800 µs   +5.1%
  potential_energy_and_grad          25 µs     25 µs  -0.0%      93.1 ms   99.6 ms   +7.0%
+ predictive_forward_sampling     702.5 ms  703.7 ms  +0.2%     208.6 ms  155.7 ms  -25.4%
  trace_seeded_model                739 µs    748 µs  +1.2%     515.5 ms  569.3 ms  +10.4%

mcmc

                             ──────── run time ───────     ────── compile time ─────
  benchmark                  baseline   this PR      Δ     baseline  this PR       Δ
────────────────────────────────────────────────────────────────────────────────────
  hmc_logistic_regression    710.8 ms  708.7 ms  -0.3%       3.34 s   3.30 s   -1.2%
  nuts_dense_mass_funnel       1.10 s    1.11 s  +1.4%       2.40 s   2.50 s   +4.2%
  nuts_eight_schools           1.09 s    1.11 s  +1.5%       2.39 s   2.50 s   +4.8%
  nuts_hierarchical_glm        4.78 s    4.78 s  -0.1%       4.82 s   5.90 s  +22.4%
  nuts_logistic_regression     1.05 s    1.06 s  +0.8%       3.12 s   3.14 s   +0.6%
  nuts_vectorized_chains       2.39 s    2.38 s  -0.4%       2.70 s   2.67 s   -1.1%

svi

                                             ──────── run time ───────     ───── compile time ─────
  benchmark                                  baseline   this PR      Δ     baseline  this PR      Δ
───────────────────────────────────────────────────────────────────────────────────────────────────
  svi_autodelta_map_logistic                 294.0 ms  299.3 ms  +1.8%       3.07 s   3.04 s  -1.0%
  svi_autodiagonalnormal_hierarchical        985.1 ms  992.7 ms  +0.8%       4.62 s   4.99 s  +7.8%
  svi_automultivariatenormal_eight_schools   716.5 ms  723.9 ms  +1.0%       3.70 s   3.98 s  +7.5%
  svi_autonormal_logistic                    732.4 ms  730.9 ms  -0.2%       3.22 s   3.45 s  +7.2%
  svi_multi_particle_elbo                      1.39 s    1.41 s  +1.4%       3.29 s   3.60 s  +9.3%
  svi_trace_mean_field_elbo                    1.26 s    1.27 s  +1.0%       5.15 s   5.12 s  -0.5%

Methodology and environment

Each benchmark is set up untimed, then called once with the JAX caches cleared and several more times warm. Run is the fastest warm call; compile is the first call minus that, i.e. the tracing, lowering and XLA compilation the warm calls did not have to pay for.

Both refs were measured on the same runner over 2 interleaved round(s), taking the best observation per benchmark. A result is called neutral when it moves less than ±5% (run) or ±25% (compile), or when the measurement itself is under 1 ms (run) / 50 ms (compile) — a shared CI runner cannot resolve changes below that. Compile time gets the looser band because it is measured once per round rather than best-of-N, and swings by roughly 20% even between two runs of identical code. A delta shown in parentheses did clear its threshold, but on a measurement below the resolution floor, so it is reported without being called a change.

	baseline	this PR
ref	master	svi-progress-bar-host-sync
commit	999d8d1f	305f8c3f
numpyro	0.21.0	0.21.0
jax	0.11.1	0.11.1
backend	cpu	cpu
python	3.14.7	3.14.7

Runner: Linux-6.17.0-1022-azure-x86_64-with-glibc2.39, 4 CPUs.

_{Produced by this benchmark run.}

[fehiepsi]

Interesting! Thanks for the enhancement!