pyro-ppl / numpyro / 32653337856

92%

push

github

Speed up posterior diagnostics (summary 1.7x, effective_sample_size 2.1x) (#2251)

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

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

14629 of 15839 relevant lines covered (92.36%)

1.69 hits per line