Repository: pyro-ppl/numpyro

Issue:
Warn on missing plate
Requested by @fehiepsi: Briefly, for models with discrete latent variables, sampling outside of plate context yields some fairly confusing stack traces during inference. Numpyro could instead warn during model validation. 

https://forum.pyro.ai/t/global-latent-variables-numpyro-crash/2313


Discussion:
Thanks, @akotlar!
Thank you Du!

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -357,6 +357,21 @@ def _guess_max_plate_nesting(model_trace):
     return max_plate_nesting
 
 
+# TODO: follow pyro.util.check_site_shape logics for more complete validation
+def _validate_model(model_trace):
+    # XXX: this validates plate statements under `enum`
+    sites = [site for site in model_trace.values()
+             if site["type"] == "sample"]
+
+    for site in sites:
+        batch_dims = len(site["fn"].batch_shape)
+        if site.get('_control_flow_done', False):
+            batch_dims = batch_dims - 1  # remove time dimension under scan
+        plate_dims = -min([0] + [frame.dim for frame in site["cond_indep_stack"]])
+        assert plate_dims >= batch_dims, \
+            "Missing plate statement for batch dimensions at site {}".format(site["name"])
+
+
 def initialize_model(rng_key, model,
                      init_strategy=init_to_uniform,
                      dynamic_args=False,
@@ -406,6 +421,7 @@ def initialize_model(rng_key, model,
 
         if not isinstance(model, enum):
             max_plate_nesting = _guess_max_plate_nesting(model_trace)
+            _validate_model(model_trace)
             model = enum(config_enumerate(model), -max_plate_nesting - 1)
 
     potential_fn, postprocess_fn = get_potential_fn(model,


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
@@ -430,3 +430,26 @@ def transition_fn(name, probs, locs, x, y):
     actual_log_joint = log_density(enum(config_enumerate(fun_model)), (data_x, data_w), {}, {})[0]
     expected_log_joint = log_density(enum(config_enumerate(model)), (data_x, data_w), {}, {})[0]
     assert_allclose(actual_log_joint, expected_log_joint)
+
+
+def test_missing_plate():
+    K, N = 3, 1000
+
+    def gmm(data):
+        mix_proportions = numpyro.sample("phi", dist.Dirichlet(jnp.ones(K)))
+        # plate/to_event is missing here
+        cluster_means = numpyro.sample("cluster_means", dist.Normal(jnp.arange(K), 1.))
+
+        with numpyro.plate("data", data.shape[0], dim=-1):
+            assignments = numpyro.sample("assignments", dist.Categorical(mix_proportions))
+            numpyro.sample("obs", dist.Normal(cluster_means[assignments], 1.), obs=data)
+
+    true_cluster_means = jnp.array([1., 5., 10.])
+    true_mix_proportions = jnp.array([0.1, 0.3, 0.6])
+    cluster_assignments = dist.Categorical(true_mix_proportions).sample(random.PRNGKey(0), (N,))
+    data = dist.Normal(true_cluster_means[cluster_assignments], 1.0).sample(random.PRNGKey(1))
+
+    nuts_kernel = NUTS(gmm)
+    mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=500)
+    with pytest.raises(AssertionError, match="Missing plate statement"):
+        mcmc.run(random.PRNGKey(2), data)