Repository: pyro-ppl/numpyro

Issue:
Only use `figure_format = svg` in website
Addresses #445.

@neerajprad It seems that logistic regression example does not work with the caching pattern. In particular, kernel.init is called with new kwargs, so we need to reconstruct `kernel._init_fn` and `kernel._sample_fn`.


Discussion:
> @neerajprad It seems that logistic regression example does not work with the caching pattern. In particular, kernel.init is called with new kwargs, so we need to reconstruct kernel._init_fn and kernel._sample_fn.

Does it throw an exception? Could you share some code snippet and elaborate on what it is doing counter to expectations?
@neerajprad For now, you can see the error with exception [here](https://nbviewer.jupyter.org/github/pyro-ppl/numpyro/blob/6537a0f81625a37e046d8bb8af934e2d4b77556d/notebooks/source/bayesian_regression.ipynb#Model-2:-Predictor---Median-Age-of-Marriage). I'll provide a simple replicate example later. :)
Oh..I was looking at the logistic regression example, not that one. Don't worry about replicating it; I'll run it and push a fix for that. 
Oppss, sorry. I meant bayesian regression. @@

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -348,6 +348,30 @@ def _find_valid_params(rng_key_):
     return init_params, is_valid
 
 
+def get_model_transforms(rng_key, model, model_args=(), model_kwargs=None):
+    model_kwargs = {} if model_kwargs is None else model_kwargs
+    seeded_model = seed(model, rng_key if rng_key.ndim == 1 else rng_key[0])
+    model_trace = trace(seeded_model).get_trace(*model_args, **model_kwargs)
+    inv_transforms = {}
+    has_transformed_dist = False
+    for k, v in model_trace.items():
+        if v['type'] == 'sample' and not v['is_observed']:
+            if v['intermediates']:
+                inv_transforms[k] = biject_to(v['fn'].base_dist.support)
+                has_transformed_dist = True
+            else:
+                inv_transforms[k] = biject_to(v['fn'].support)
+        elif v['type'] == 'param':
+            constraint = v['kwargs'].pop('constraint', real)
+            transform = biject_to(constraint)
+            if isinstance(transform, ComposeTransform):
+                inv_transforms[k] = transform.parts[0]
+                has_transformed_dist = True
+            else:
+                inv_transforms[k] = transform
+    return inv_transforms, has_transformed_dist
+
+
 def get_potential_fn(rng_key, model, dynamic_args=False, model_args=(), model_kwargs=None):
     """
     (EXPERIMENTAL INTERFACE) Given a model with Pyro primitives, returns a
@@ -370,37 +394,19 @@ def get_potential_fn(rng_key, model, dynamic_args=False, model_args=(), model_kw
         to constrain unconstrained samples (e.g. those returned by HMC)
         to values that lie within the site's support.
     """
-    model_kwargs = {} if model_kwargs is None else model_kwargs
-    seeded_model = seed(model, rng_key if rng_key.ndim == 1 else rng_key[0])
-    model_trace = trace(seeded_model).get_trace(*model_args, **model_kwargs)
-    inv_transforms = {}
-    has_transformed_dist = False
-    for k, v in model_trace.items():
-        if v['type'] == 'sample' and not v['is_observed']:
-            if v['intermediates']:
-                inv_transforms[k] = biject_to(v['fn'].base_dist.support)
-                has_transformed_dist = True
-            else:
-                inv_transforms[k] = biject_to(v['fn'].support)
-        elif v['type'] == 'param':
-            constraint = v['kwargs'].pop('constraint', real)
-            transform = biject_to(constraint)
-            if isinstance(transform, ComposeTransform):
-                inv_transforms[k] = transform.parts[0]
-                has_transformed_dist = True
-            else:
-                inv_transforms[k] = transform
-
     if dynamic_args:
         def potential_fn(*args, **kwargs):
+            inv_transforms, has_transformed_dist = get_model_transforms(rng_key, model, args, kwargs)
             return jax.partial(potential_energy, model, inv_transforms, args, kwargs)
-        if has_transformed_dist:
-            def constrain_fun(*args, **kwargs):
+
+        def constrain_fun(*args, **kwargs):
+            inv_transforms, has_transformed_dist = get_model_transforms(rng_key, model, args, kwargs)
+            if has_transformed_dist:
                 return jax.partial(constrain_fn, model, inv_transforms, args, kwargs)
-        else:
-            def constrain_fun(*args, **kwargs):
+            else:
                 return jax.partial(transform_fn, inv_transforms)
     else:
+        inv_transforms, has_transformed_dist = get_model_transforms(rng_key, model, model_args, model_kwargs)
         potential_fn = jax.partial(potential_energy, model, inv_transforms, model_args, model_kwargs)
         if has_transformed_dist:
             constrain_fun = jax.partial(constrain_fn, model, inv_transforms, model_args, model_kwargs)


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -491,3 +491,32 @@ def model(y_obs):
     # Re-run on new data - should be much faster.
     mcmc.run(random.PRNGKey(32), y2)
     assert_allclose(mcmc.get_samples()['mu'].mean(), -3., atol=0.1)
+
+
+@pytest.mark.parametrize('jit_args', [False, True])
+def test_model_with_multiple_exec_paths(jit_args):
+    def model(a=None, b=None, z=None):
+        int_term = numpyro.sample('a', dist.Normal(0., 0.2))
+        x_term, y_term = 0., 0.
+        if a is not None:
+            x = numpyro.sample('x', dist.HalfNormal(0.5))
+            x_term = a * x
+        if b is not None:
+            y = numpyro.sample('y', dist.HalfNormal(0.5))
+            y_term = b * y
+        sigma = numpyro.sample('sigma', dist.Exponential(1.))
+        mu = int_term + x_term + y_term
+        numpyro.sample('obs', dist.Normal(mu, sigma), obs=z)
+
+    a = np.exp(onp.random.randn(10))
+    b = np.exp(onp.random.randn(10))
+    z = onp.random.randn(10)
+
+    # Run MCMC on zero observations.
+    kernel = NUTS(model)
+    mcmc = MCMC(kernel, 20, jit_model_args=jit_args)
+    mcmc.run(random.PRNGKey(0), a, b=None, z=z)
+    mcmc.num_samples = 10
+    mcmc.run(random.PRNGKey(1), a, b=None, z=z, reuse_warmup_state=True)
+    mcmc.run(random.PRNGKey(2), a=None, b=b, z=z)
+    mcmc.run(random.PRNGKey(3), a=a, b=b, z=z)