Repository: pyro-ppl/numpyro

Issue:
Raise warning for the automatic enumerate behavior
Following a thread on NumPyro forum about inference with discrete latent variable in hmm, here is the github issue as requested: 
In the hmm example, the discrete latent variable are marginalized out automatically. 
To avoid confusion, in future versions, the "automatic" behavior could be removed, and the users would provide infer={"enumerate": "parallel"} explicitly.
Thank you.


Discussion:

Resolving patch:
diff --git a/examples/annotation.py b/examples/annotation.py
--- a/examples/annotation.py
+++ b/examples/annotation.py
@@ -123,7 +123,7 @@ def multinomial(annotations):
     pi = numpyro.sample("pi", dist.Dirichlet(jnp.ones(num_classes)))
 
     with numpyro.plate("item", num_items, dim=-2):
-        c = numpyro.sample("c", dist.Categorical(pi))
+        c = numpyro.sample("c", dist.Categorical(pi), infer={"enumerate": "parallel"})
 
         with numpyro.plate("position", num_positions):
             numpyro.sample("y", dist.Categorical(zeta[c]), obs=annotations)
@@ -144,7 +144,7 @@ def dawid_skene(positions, annotations):
     pi = numpyro.sample("pi", dist.Dirichlet(jnp.ones(num_classes)))
 
     with numpyro.plate("item", num_items, dim=-2):
-        c = numpyro.sample("c", dist.Categorical(pi))
+        c = numpyro.sample("c", dist.Categorical(pi), infer={"enumerate": "parallel"})
 
         # here we use Vindex to allow broadcasting for the second index `c`
         # ref: http://num.pyro.ai/en/latest/utilities.html#numpyro.contrib.indexing.vindex
@@ -167,10 +167,18 @@ def mace(positions, annotations):
         theta = numpyro.sample("theta", dist.Beta(0.5, 0.5))
 
     with numpyro.plate("item", num_items, dim=-2):
-        c = numpyro.sample("c", dist.DiscreteUniform(0, num_classes - 1))
+        c = numpyro.sample(
+            "c",
+            dist.DiscreteUniform(0, num_classes - 1),
+            infer={"enumerate": "parallel"},
+        )
 
         with numpyro.plate("position", num_positions):
-            s = numpyro.sample("s", dist.Bernoulli(1 - theta[positions]))
+            s = numpyro.sample(
+                "s",
+                dist.Bernoulli(1 - theta[positions]),
+                infer={"enumerate": "parallel"},
+            )
             probs = jnp.where(
                 s[..., None] == 0, nn.one_hot(c, num_classes), epsilon[positions]
             )
@@ -207,7 +215,7 @@ def hierarchical_dawid_skene(positions, annotations):
     pi = numpyro.sample("pi", dist.Dirichlet(jnp.ones(num_classes)))
 
     with numpyro.plate("item", num_items, dim=-2):
-        c = numpyro.sample("c", dist.Categorical(pi))
+        c = numpyro.sample("c", dist.Categorical(pi), infer={"enumerate": "parallel"})
 
         with numpyro.plate("position", num_positions):
             logits = Vindex(beta)[positions, c, :]
@@ -232,7 +240,7 @@ def item_difficulty(annotations):
     pi = numpyro.sample("pi", dist.Dirichlet(jnp.ones(num_classes)))
 
     with numpyro.plate("item", num_items, dim=-2):
-        c = numpyro.sample("c", dist.Categorical(pi))
+        c = numpyro.sample("c", dist.Categorical(pi), infer={"enumerate": "parallel"})
 
         with handlers.reparam(config={"theta": LocScaleReparam(0)}):
             theta = numpyro.sample("theta", dist.Normal(eta[c], chi[c]).to_event(1))
@@ -270,7 +278,7 @@ def logistic_random_effects(positions, annotations):
     pi = numpyro.sample("pi", dist.Dirichlet(jnp.ones(num_classes)))
 
     with numpyro.plate("item", num_items, dim=-2):
-        c = numpyro.sample("c", dist.Categorical(pi))
+        c = numpyro.sample("c", dist.Categorical(pi), infer={"enumerate": "parallel"})
 
         with handlers.reparam(config={"theta": LocScaleReparam(0)}):
             theta = numpyro.sample("theta", dist.Normal(0, chi[c]).to_event(1))
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -402,11 +402,28 @@ def _get_model_transforms(model, model_args=(), model_kwargs=None):
     for k, v in model_trace.items():
         if v["type"] == "sample" and not v["is_observed"]:
             if v["fn"].support.is_discrete:
+                enum_type = v["infer"].get("enumerate")
+                if enum_type is not None and (enum_type != "parallel"):
+                    raise RuntimeError(
+                        "This algorithm might only work for discrete sites with"
+                        f" enumerate marked 'parallel'. But the site {k} is marked"
+                        f" as '{enum_type}'."
+                    )
                 has_enumerate_support = True
                 if not v["fn"].has_enumerate_support:
+                    dist_name = type(v["fn"]).__name__
                     raise RuntimeError(
-                        "MCMC only supports continuous sites or discrete sites "
-                        f"with enumerate support, but got {type(v['fn']).__name__}."
+                        "This algorithm might only work for discrete sites with"
+                        f" enumerate support. But the {dist_name} distribution at"
+                        f" site {k} does not have enumerate support."
+                    )
+                if enum_type is None:
+                    warnings.warn(
+                        "Some algorithms will automatically enumerate the discrete"
+                        f" latent site {k} of your model. In the future,"
+                        " enumerated sites need to be marked with"
+                        " `infer={'enumerate': 'parallel'}`.",
+                        FutureWarning,
                     )
             else:
                 support = v["fn"].support


Test patch:
diff --git a/test/contrib/test_funsor.py b/test/contrib/test_funsor.py
--- a/test/contrib/test_funsor.py
+++ b/test/contrib/test_funsor.py
@@ -26,6 +26,7 @@
 def test_gaussian_mixture_model():
     K, N = 3, 1000
 
+    @config_enumerate
     def gmm(data):
         mix_proportions = numpyro.sample("phi", dist.Dirichlet(jnp.ones(K)))
         with numpyro.plate("num_clusters", K, dim=-1):
@@ -60,6 +61,7 @@ def gmm(data):
 
 
 def test_bernoulli_latent_model():
+    @config_enumerate
     def model(data):
         y_prob = numpyro.sample("y_prob", dist.Beta(1.0, 1.0))
         with numpyro.plate("data", data.shape[0]):
@@ -81,6 +83,7 @@ def model(data):
 
 
 def test_change_point():
+    @config_enumerate
     def model(count_data):
         n_count_data = count_data.shape[0]
         alpha = 1 / jnp.mean(count_data.astype(np.float32))
@@ -93,84 +96,13 @@ def model(count_data):
         with numpyro.plate("data", n_count_data):
             numpyro.sample("obs", dist.Poisson(lambda_), obs=count_data)
 
-    count_data = jnp.array(
-        [
-            13,
-            24,
-            8,
-            24,
-            7,
-            35,
-            14,
-            11,
-            15,
-            11,
-            22,
-            22,
-            11,
-            57,
-            11,
-            19,
-            29,
-            6,
-            19,
-            12,
-            22,
-            12,
-            18,
-            72,
-            32,
-            9,
-            7,
-            13,
-            19,
-            23,
-            27,
-            20,
-            6,
-            17,
-            13,
-            10,
-            14,
-            6,
-            16,
-            15,
-            7,
-            2,
-            15,
-            15,
-            19,
-            70,
-            49,
-            7,
-            53,
-            22,
-            21,
-            31,
-            19,
-            11,
-            1,
-            20,
-            12,
-            35,
-            17,
-            23,
-            17,
-            4,
-            2,
-            31,
-            30,
-            13,
-            27,
-            0,
-            39,
-            37,
-            5,
-            14,
-            13,
-            22,
-        ]
-    )
+    # fmt: off
+    count_data = jnp.array([
+        13, 24, 8, 24, 7, 35, 14, 11, 15, 11, 22, 22, 11, 57, 11, 19, 29, 6, 19, 12, 22,
+        12, 18, 72, 32, 9, 7, 13, 19, 23, 27, 20, 6, 17, 13, 10, 14, 6, 16, 15, 7, 2,
+        15, 15, 19, 70, 49, 7, 53, 22, 21, 31, 19, 11, 1, 20, 12, 35, 17, 23, 17, 4, 2,
+        31, 30, 13, 27, 0, 39, 37, 5, 14, 13, 22])
+    # fmt: on
 
     kernel = NUTS(model)
     mcmc = MCMC(kernel, num_warmup=500, num_samples=500)
@@ -184,6 +116,7 @@ def test_gaussian_hmm():
     dim = 4
     num_steps = 10
 
+    @config_enumerate
     def model(data):
         with numpyro.plate("states", dim):
             transition = numpyro.sample("transition", dist.Dirichlet(jnp.ones(dim)))
@@ -586,6 +519,7 @@ def transition_fn(carry, y):
 def test_missing_plate(monkeypatch):
     K, N = 3, 1000
 
+    @config_enumerate
     def gmm(data):
         mix_proportions = numpyro.sample("phi", dist.Dirichlet(jnp.ones(K)))
         # plate/to_event is missing here
diff --git a/test/contrib/test_infer_discrete.py b/test/contrib/test_infer_discrete.py
--- a/test/contrib/test_infer_discrete.py
+++ b/test/contrib/test_infer_discrete.py
@@ -405,7 +405,7 @@ def model2():
 @pytest.mark.parametrize("model", [model_zzxx, model2])
 @pytest.mark.parametrize("temperature", [0, 1])
 def test_mcmc_model_side_enumeration(model, temperature):
-    mcmc = infer.MCMC(infer.NUTS(model), num_warmup=0, num_samples=1)
+    mcmc = infer.MCMC(infer.NUTS(config_enumerate(model)), num_warmup=0, num_samples=1)
     mcmc.run(random.PRNGKey(0))
     mcmc_data = {
         k: v[0] for k, v in mcmc.get_samples().items() if k in ["loc", "scale"]
diff --git a/test/infer/test_autoguide.py b/test/infer/test_autoguide.py
--- a/test/infer/test_autoguide.py
+++ b/test/infer/test_autoguide.py
@@ -562,6 +562,7 @@ def model():
         init_to_uniform,
     ],
 )
+@pytest.mark.filterwarnings("ignore:.*enumerate.*:FutureWarning")
 def test_discrete_helpful_error(auto_class, init_loc_fn):
     def model():
         p = numpyro.sample("p", dist.Beta(2.0, 2.0))
diff --git a/test/infer/test_hmc_gibbs.py b/test/infer/test_hmc_gibbs.py
--- a/test/infer/test_hmc_gibbs.py
+++ b/test/infer/test_hmc_gibbs.py
@@ -281,7 +281,7 @@ def test_block_update_partitioning(num_blocks):
 
 def test_enum_subsample_smoke():
     def model(data):
-        x = numpyro.sample("x", dist.Bernoulli(0.5))
+        x = numpyro.sample("x", dist.Bernoulli(0.5), infer={"enumerate": "parallel"})
         with numpyro.plate("N", data.shape[0], subsample_size=100, dim=-1):
             batch = numpyro.subsample(data, event_dim=0)
             numpyro.sample("obs", dist.Normal(x, 1), obs=batch)
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -1060,3 +1060,12 @@ def model():
     mcmc = MCMC(NUTS(subs_model), num_warmup=10, num_samples=10)
     with pytest.warns(UserWarning, match="skipping initialization"):
         mcmc.run(random.PRNGKey(1))
+
+
+def test_discrete_site_without_infer_enumerate():
+    def model():
+        numpyro.sample("x", dist.Bernoulli(0.5))
+
+    mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
+    with pytest.warns(FutureWarning, match="enumerated sites"):
+        mcmc.run(random.PRNGKey(0))