Numpyro: Investigate discrepancy in effective sample size computation

[neerajprad]

The effective sample size as reported in NumPyro / Pyro seems to be lower than the one reported by Stan. Details in https://github.com/pyro-ppl/brmp/pull/25.

[neerajprad]

Taking the discussion over from https://github.com/pyro-ppl/brmp/pull/25:

Update2: After making many mistakes, it seems to me that both Stan and NumPyro have the same way to compute n_eff. I will look further to see why NumPyro gives smaller n_eff for the topic model.

I think one way to validate that the computation is the same is to pass the samples collected from brms (or pystan) to our diagnostics package. If the reported n_eff matches the one from brms, then there is probably some efficiency issue with our sampler, at least for this model.

[null-a]

one way to validate that the computation is the same is to pass the samples collected from brms (or pystan) to our diagnostics package

I just did this (using post warmup samples only) for the example I have in pyro-ppl/brmp#25 and found that the output of numpyro.diagnostics.effective_sample_size matches the n_eff reported by brms.

[neerajprad]

Marking this as fixed. Please feel free to reopen if you notice this issue on other models. I think @fehiepsi's fix seems to work well on a variety of models that we looked at.