Speed up posterior diagnostics (summary 1.7x, effective_sample_size 2.1x) - #2251
Conversation
Three changes to numpyro/diagnostics.py: - Make autocovariance the FFT primitive and derive autocorrelation from it by lag-0 normalization, instead of computing autocorrelation first and multiplying the variance back in. This removes two full-array passes, a separate variance pass, and a redundant float64 astype copy from the effective_sample_size path. - Use scipy.fft with workers=-1 (scipy is always present via jax; fall back to np.fft if unavailable) so batched FFTs use all cores. - In summary(), sort the draws once and reuse the sorted array for both the median and the hpdi bounds; previously np.median partitioned the same array that hpdi then fully sorted. summary() is 1.6-1.7x faster and effective_sample_size 2.1x faster on realistic posterior sizes; float64 results agree with master to rtol 1e-9 (float32 inputs now normalize in double precision, removing float32-level round-off from n_eff).
|
Thanks @kyo219! Can you share the code you used for benchmarks. |
Benchmark reportthis PR - run time: 1 slower, 0 faster
compile time: unchanged across 32 benchmarksSignificant changes (1) ─────── run time ─────── ────── compile time ─────
benchmark baseline this PR Δ baseline this PR Δ
────────────────────────────────────────────────────────────────────────
- normal_sample 24.9 ms 26.3 ms +5.8% 212.7 ms 216.0 ms +1.6%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
|
| baseline | this PR | |
|---|---|---|
| ref | master |
diagnostics-speedup |
| commit | 999d8d1f |
3304780d |
| 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.
|
Sure! Script below (Apple M2, 8 cores, CPU; best of 3 after warmup). Agreement was checked separately by dumping every statistic ( import time
import numpy as np
from numpyro.diagnostics import summary, effective_sample_size
def timeit(fn, n=3):
fn()
ts = []
for _ in range(n):
t0 = time.perf_counter(); fn(); ts.append(time.perf_counter() - t0)
return min(ts)
rng = np.random.default_rng(0)
small = {"a": rng.standard_normal((4, 2000, 100)), "b": rng.standard_normal((4, 2000, 50, 10))}
big = {"c": rng.standard_normal((4, 5000, 2000))}
print(f"summary small: {timeit(lambda: summary(small)):.3f}s")
print(f"summary big: {timeit(lambda: summary(big)):.3f}s")
print(f"ess: {timeit(lambda: effective_sample_size(small['b'])):.3f}s") |
2 participants
Changes made
All in
numpyro/diagnostics.py:autocovarianceis now the FFT primitive,autocorrelationderives from it. Previouslyautocorrelationnormalized the raw FFT autocovariance by its lag-0 value, andautocovariancethen multiplied the variance back in (autocorrelation(x) * x.var()) — two extra full-array passes plus a separate variance pass, and a redundant.astype(np.float64)copy (the FFT output is already float64).effective_sample_sizeconsumesautocovariance, so its hot path drops all of these.autocorrelationstill returns the same values, now computed asautocov / autocov[lag 0](the lag-0 autocovariance is identical for the biased and unbiased estimators, which is what the old normalization relied on implicitly).scipy.fftwithworkers=-1, with annp.fftfallback. scipy is not a direct numpyro dependency but is guaranteed transitively via jax, and numpyro already uses it in several lazy imports (e.g.distributions/continuous.py,infer/calibration.py); the test suite also importsscipy.fftpackintest_diagnostics.py.np.fftis single-threaded;scipy.fftparallelizes across the batch dimensions, which is exactly the shape of the ESS workload (one FFT per chain per parameter column). Sincescipy.fftcomputes in the input precision (unlikenp.fft, which always promotes), the centered signal is explicitly cast to float64 to keep the double-precision behavior.summary()sorts each parameter's draws once and reuses the sorted array for both the median and the hpdi bounds. Previouslynp.medianpartitioned the very array thathpdithen fully sorted. The sliding-window part ofhpdiis factored into a private_hpdi_of_sortedhelper; the publichpdiis unchanged. The median-from-sorted matchesnp.mediansemantics including NaN propagation (sorting places NaNs last, so checking the last draw suffices).Behavioral notes: float64 results agree with master to rtol 1e-9. For float32 inputs, results differ at the float32-round-off level (relative ~1e-7) because the old code normalized with a float32
x.var()while the new code stays in float64 throughout — the new values are the more precise ones.autocovarianceof a constant series now returns 0 rather than NaN (the old NaN came from the 0/0 lag-0 normalization that the variance multiplication couldn't undo);effective_sample_sizeandautocorrelationstill return NaN for that case as before.Benchmarks
Apple M2 (24 GB, 8 cores), CPU, numpy 2.x/scipy via jax 0.10.2, Python 3.11. Best of 3 after warmup.
summary()on {a: (4, 2000, 100), b: (4, 2000, 50, 10)}summary()on (4, 5000, 2000)effective_sample_sizeon (4, 2000, 50, 10)hpdistandalone (8000, 500)This is the
mcmc.print_summary()path, so every MCMC run with the default summary printout benefits.Links to related issues/PRs
None.
Tests
test_summary_median_matches_numpy(odd and even draw counts) andtest_summary_median_propagates_nanintest/test_diagnostics.py— lock in that the shared-sort median is exactlynp.median, including NaN propagation.pytest test/test_diagnostics.py— 28 passed (autocorrelation/autocovariance/ESS/hpdi/gelman-rubin values are asserted against fixed references there).autocorrelation/autocovariance(bias=True/False),effective_sample_size,hpdi, and everysummarystatistic: float64 agrees to rtol 1e-9; float32 differs only at round-off as described above.ruff check/ruff format --check/ty checkclean.Dependencies
None (scipy is used only behind a try/except fallback).