Repository: pyro-ppl/numpyro

Issue:
Improvements to help benchmark numpyro code
 - [x] Ability to reuse the compiled model + adapted state to rerun sampling from different seeds.
 - [x] Separating `summary` statistics from `print_summary`
 - [x] Add support for using x64. 
 - [ ] Fix perf regression when progress bar is disabled.
cc. @fehiepsi - feel free to add to this.


Discussion:

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -6,7 +6,7 @@
 import os
 import warnings
 
-from jax import jit, lax, partial, pmap, random, vmap, device_get
+from jax import jit, lax, partial, pmap, random, vmap, device_get, device_put
 from jax.core import Tracer
 from jax.dtypes import canonicalize_dtype
 from jax.flatten_util import ravel_pytree
@@ -26,7 +26,7 @@
     warmup_adapter
 )
 from numpyro.infer.util import init_to_uniform, get_potential_fn, find_valid_initial_params
-from numpyro.util import cond, copy_docs_from, fori_collect, fori_loop, identity, not_jax_tracer
+from numpyro.util import cond, copy_docs_from, fori_collect, fori_loop, identity, not_jax_tracer, cached_by
 
 HMCState = namedtuple('HMCState', ['i', 'z', 'z_grad', 'potential_energy', 'energy', 'num_steps', 'accept_prob',
                                    'mean_accept_prob', 'diverging', 'adapt_state', 'rng_key'])
@@ -254,7 +254,7 @@ def init_kernel(init_params,
         energy = kinetic_fn(wa_state.inverse_mass_matrix, vv_state.r)
         hmc_state = HMCState(0, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
                              0, 0., 0., False, wa_state, rng_key_hmc)
-        return hmc_state
+        return device_put(hmc_state)
 
     def _hmc_next(step_size, inverse_mass_matrix, vv_state,
                   model_args, model_kwargs, rng_key):
@@ -635,6 +635,14 @@ def _sample_fn_nojit_args(state, sampler, args, kwargs):
     return sampler.sample(state[0], args, kwargs),
 
 
+def _collect_fn(collect_fields):
+    @cached_by(_collect_fn, collect_fields)
+    def collect(x):
+        return attrgetter(*collect_fields)(x[0])
+
+    return collect
+
+
 class MCMC(object):
     """
     Provides access to Markov Chain Monte Carlo inference algorithms in NumPyro.
@@ -732,20 +740,19 @@ def _single_chain_mcmc(self, rng_key, init_params, args, kwargs,
             if self._warmup_state is None:
                 raise ValueError('No `init_state` found; warmup results cannot be reused.')
             num_warmup = 0
-            init_state = self._warmup_state
+            init_state = self._warmup_state._replace(rng_key=rng_key)
         else:
             init_state = self.sampler.init(rng_key, self.num_warmup, init_params,
                                            model_args=args, model_kwargs=kwargs)
         if self.constrain_fn is None:
             self.constrain_fn = self.sampler.constrain_fn(args, kwargs)
-        collect_fn = attrgetter(*collect_fields)
         lower = 0 if collect_warmup else num_warmup
         diagnostics = lambda x: get_diagnostics_str(x[0]) if rng_key.ndim == 1 else None   # noqa: E731
         init_val = (init_state, args, kwargs) if self._jit_model_args else (init_state,)
         collect_vals = fori_collect(lower, num_warmup + self.num_samples,
                                     self._get_cached_fn(),
                                     init_val,
-                                    transform=lambda x: collect_fn(x[0]),  # noqa: E731
+                                    transform=_collect_fn(collect_fields),  # noqa: E731
                                     progbar=self.progress_bar,
                                     return_init_state=True,
                                     progbar_desc=functools.partial(get_progbar_desc_str, num_warmup),
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -142,6 +142,19 @@ def identity(x):
     return x
 
 
+def cached_by(outer_fn, *keys):
+    outer_fn._cache = getattr(outer_fn, '_cache', {})
+
+    def _wrapped(fn):
+        fn_cache = outer_fn._cache
+        if keys in fn_cache:
+            return fn_cache[keys]
+        fn_cache[keys] = fn
+        return fn
+
+    return _wrapped
+
+
 def fori_collect(lower, upper, body_fun, init_val, transform=identity,
                  progbar=True, return_init_state=False, **progbar_opts):
     """
@@ -179,17 +192,18 @@ def fori_collect(lower, upper, body_fun, init_val, transform=identity,
     if not progbar:
         collection = np.zeros((upper - lower,) + init_val_flat.shape)
 
+        @cached_by(fori_collect, body_fun, transform)
         def _body_fn(i, vals):
-            val, collection, start_state = vals
-            val = jit(body_fun)(val)
-            start_state = lax.cond(i < lower,
+            val, collection, start_state, lower_idx = vals
+            val = body_fun(val)
+            start_state = lax.cond(i < lower_idx,
                                    val, lambda x: x,
                                    start_state, lambda x: x)
-            i = np.where(i >= lower, i - lower, 0)
-            collection = ops.index_update(collection, i, ravel_fn(val))
-            return val, collection, start_state
+            i = np.where(i >= lower_idx, i - lower_idx, 0)
+            collection = ops.index_update(collection, i, ravel_pytree(transform(val))[0])
+            return val, collection, start_state, lower_idx
 
-        _, collection, start_state = fori_loop(0, upper, _body_fn, (init_val, collection, init_val))
+        _, collection, start_state, _ = fori_loop(0, upper, _body_fn, (init_val, collection, init_val, lower))
     else:
         diagnostics_fn = progbar_opts.pop('diagnostics_fn', None)
         progbar_desc = progbar_opts.pop('progbar_desc', lambda x: '')


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -1,6 +1,7 @@
 import os
 
 import numpy as onp
+from jax.test_util import check_close
 from numpy.testing import assert_allclose
 import pytest
 
@@ -248,6 +249,31 @@ def model(data):
     assert_allclose(np.mean(samples['std']), true_std, rtol=0.05)
 
 
+def test_mcmc_progbar():
+    true_mean, true_std = 1., 2.
+    num_warmup, num_samples = 10, 10
+
+    def model(data):
+        mean = numpyro.param('mean', 0.)
+        std = numpyro.param('std', 1., constraint=constraints.positive)
+        return numpyro.sample('obs', dist.Normal(mean, std), obs=data)
+
+    data = dist.Normal(true_mean, true_std).sample(random.PRNGKey(1), (2000,))
+    kernel = NUTS(model=model)
+    mcmc = MCMC(kernel, num_warmup)
+    mcmc.run(random.PRNGKey(2), data)
+    mcmc.num_samples = num_samples
+    mcmc.run(random.PRNGKey(3), data, reuse_warmup_state=True)
+    mcmc1 = MCMC(kernel, num_warmup, num_samples, progress_bar=False)
+    mcmc1.run(random.PRNGKey(2), data)
+
+    with pytest.raises(AssertionError):
+        check_close(mcmc1.get_samples(), mcmc.get_samples(), atol=1e-3, rtol=0.01)
+    mcmc1.run(random.PRNGKey(3), data, reuse_warmup_state=True)
+    check_close(mcmc1.get_samples(), mcmc.get_samples(), atol=1e-3, rtol=0.01)
+    check_close(mcmc1._warmup_state, mcmc._warmup_state, atol=1e-3, rtol=0.01)
+
+
 @pytest.mark.parametrize('kernel_cls', [HMC, NUTS])
 @pytest.mark.parametrize('adapt_step_size', [True, False])
 def test_diverging(kernel_cls, adapt_step_size):