Repository: pyro-ppl/numpyro (python)

Issue:
HMCGibbs with chain_method=”vectorized”
I am trying to use HMCGibbs sampling with more than one chain using chain_method=“vectorized”, but there appears to be some problem with splitting the random keys.

Consider this toy example that I copied from the numpyro documentation, where I only changed the chain_method and the number of chains:

```
from jax import random
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS, HMCGibbs

def model():
    x = numpyro.sample("x", dist.Normal(0.0, 2.0))
    y = numpyro.sample("y", dist.Normal(0.0, 2.0))
    numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0]))

def gibbs_fn(rng_key, gibbs_sites, hmc_sites):
    y = hmc_sites['y']
    new_x = dist.Normal(0.8 * (1-y), jnp.sqrt(0.8)).sample(rng_key)
    return {'x': new_x}

hmc_kernel = NUTS(model)
kernel = HMCGibbs(hmc_kernel, gibbs_fn=gibbs_fn, gibbs_sites=['x'])
mcmc = MCMC(kernel, num_warmup=100, num_chains=2, num_samples=100, progress_bar=False, chain_method='vectorized',)
mcmc.run(random.PRNGKey(0))
mcmc.print_summary()
```


I find that I get the following Error when running the above Code:
`TypeError: split accepts a single key, but was given a key array of shape (2,) != (). Use jax.vmap for batching.`

Is there a way to make the vectorize option available for HMCGibbs sampling?


Pull request description:
Allow for more general chain_method in MCMC
Fixes #1725

Currently, we only support "vectorized" chain method in HMC/NUTS. This PR allows chain_method to be a callable, so that users can use vectorized method via `chain_method=jax.vmap`.

I found this is simpler than supporting vmap explicitly for each kernel.

Patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
index 6e13aeb7..11fdde45 100644
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -817,7 +817,7 @@ class substitute(Messenger):
             return
 
         if self.data is not None:
-            value = self.data.get(msg["name"])
+            value = self.data.get(msg.get("name"))
         else:
             value = self.substitute_fn(msg)
 
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
index ad016825..e18b35bc 100644
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -277,7 +277,8 @@ class MCMC(object):
         sample values returned from the sampler to constrained values that lie within the support
         of the sample sites. Additionally, this is used to return values at deterministic sites in
         the model.
-    :param str chain_method: One of 'parallel' (default), 'sequential', 'vectorized'. The method
+    :param str chain_method: A callable jax transform like `jax.vmap` or one of
+        'parallel' (default), 'sequential', 'vectorized'. The method
         'parallel' is used to execute the drawing process in parallel on XLA devices (CPUs/GPUs/TPUs),
         If there are not enough devices for 'parallel', we fall back to 'sequential' method to draw
         chains sequentially. 'vectorized' method is an experimental feature which vectorizes the
@@ -340,7 +341,11 @@ class MCMC(object):
             raise ValueError("thinning must be a positive integer")
         self.thinning = thinning
         self.postprocess_fn = postprocess_fn
-        if chain_method not in ["parallel", "vectorized", "sequential"]:
+        if not callable(chain_method) and chain_method not in [
+            "parallel",
+            "vectorized",
+            "sequential",
+        ]:
             raise ValueError(
                 "Only supporting the following methods to draw chains:"
                 ' "sequential", "parallel", or "vectorized"'
@@ -471,7 +476,9 @@ class MCMC(object):
             collection_size=collection_size,
             progbar_desc=partial(_get_progbar_desc_str, lower_idx, phase),
             diagnostics_fn=diagnostics,
-            num_chains=self.num_chains if self.chain_method == "parallel" else 1,
+            num_chains=self.num_chains
+            if (callable(self.chain_method) or self.chain_method == "parallel")
+            else 1,
         )
         states, last_val = collect_vals
         # Get first argument of type `HMCState`
@@ -679,6 +686,8 @@ class MCMC(object):
                 states, last_state = _laxmap(partial_map_fn, map_args)
             elif self.chain_method == "parallel":
                 states, last_state = pmap(partial_map_fn)(map_args)
+            elif callable(self.chain_method):
+                states, last_state = self.chain_method(partial_map_fn)(map_args)
             else:
                 assert self.chain_method == "vectorized"
                 states, last_state = partial_map_fn(map_args)
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
index 3775893a..56c4d2c4 100644
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -21,6 +21,7 @@ from numpyro.distributions.transforms import biject_to
 from numpyro.distributions.util import is_identically_one, sum_rightmost
 from numpyro.handlers import condition, replay, seed, substitute, trace
 from numpyro.infer.initialization import init_to_uniform, init_to_value
+from numpyro.primitives import Messenger
 from numpyro.util import (
     _validate_model,
     find_stack_level,
@@ -46,6 +47,12 @@ ModelInfo = namedtuple(
 ParamInfo = namedtuple("ParamInfo", ["z", "potential_energy", "z_grad"])
 
 
+class _substitute_default_key(Messenger):
+    def process_message(self, msg):
+        if msg["type"] == "prng_key" and msg["value"] is None:
+            msg["value"] = random.PRNGKey(0)
+
+
 def log_density(model, model_args, model_kwargs, params):
     """
     (EXPERIMENTAL INTERFACE) Computes log of joint density for the model given
@@ -660,9 +667,12 @@ def initialize_model(
         data={
             k: site["value"]
             for k, site in model_trace.items()
-            if site["type"] in ["param"]
+            if site["type"] in ["param", "mutable"]
         },
     )
+
+    model = _substitute_default_key(model)
+
     constrained_values = {
         k: v["value"]
         for k, v in model_trace.items()


Test patch:
diff --git a/test/contrib/test_module.py b/test/contrib/test_module.py
index d3f27f17..a1342507 100644
--- a/test/contrib/test_module.py
+++ b/test/contrib/test_module.py
@@ -21,7 +21,8 @@ from numpyro.contrib.module import (
     random_haiku_module,
 )
 import numpyro.distributions as dist
-from numpyro.infer import MCMC, NUTS
+from numpyro.infer import MCMC, NUTS, SVI, Trace_ELBO
+from numpyro.infer.autoguide import AutoDelta
 
 pytestmark = pytest.mark.filterwarnings(
     "ignore:jax.tree_.+ is deprecated:FutureWarning"
@@ -256,6 +257,11 @@ def test_haiku_state_dropout_smoke(dropout, batchnorm):
     else:
         assert set(tr.keys()) == {"nn$params", "x", "y"}
 
+    # test svi
+    guide = AutoDelta(model)
+    svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO())
+    svi.run(random.PRNGKey(100), 10)
+
 
 @pytest.mark.parametrize("dropout", [True, False])
 @pytest.mark.parametrize("batchnorm", [True, False])
@@ -300,3 +306,8 @@ def test_flax_state_dropout_smoke(dropout, batchnorm):
         assert tr["nn$state"]["type"] == "mutable"
     else:
         assert set(tr.keys()) == {"nn$params", "x", "y"}
+
+    # test svi
+    guide = AutoDelta(model)
+    svi = SVI(model, guide, numpyro.optim.Adam(0.01), Trace_ELBO())
+    svi.run(random.PRNGKey(100), 10)
diff --git a/test/infer/test_hmc_gibbs.py b/test/infer/test_hmc_gibbs.py
index 3cb7f02a..5caef6f3 100644
--- a/test/infer/test_hmc_gibbs.py
+++ b/test/infer/test_hmc_gibbs.py
@@ -462,3 +462,24 @@ def test_hmcecs_multiple_plates():
     kernel = HMCECS(NUTS(model), proxy=proxy_fn)
     mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
     mcmc.run(random.PRNGKey(0), data)
+
+
+def test_callable_chain_method():
+    def model():
+        x = numpyro.sample("x", dist.Normal(0.0, 2.0))
+        y = numpyro.sample("y", dist.Normal(0.0, 2.0))
+        numpyro.sample("obs", dist.Normal(x + y, 1.0), obs=jnp.array([1.0]))
+
+    def gibbs_fn(rng_key, gibbs_sites, hmc_sites):
+        y = hmc_sites["y"]
+        new_x = dist.Normal(0.8 * (1 - y), jnp.sqrt(0.8)).sample(rng_key)
+        return {"x": new_x}
+
+    hmc_kernel = NUTS(model)
+    kernel = HMCGibbs(hmc_kernel, gibbs_fn=gibbs_fn, gibbs_sites=["x"])
+    mcmc = MCMC(
+        kernel, num_warmup=100, num_chains=2, num_samples=100, chain_method=vmap
+    )
+    mcmc.run(random.PRNGKey(0))
+    samples = mcmc.get_samples()
+    assert set(samples.keys()) == {"x", "y"}