Add a warning for missing plate statements in the model (#1245)

* Fix missing plate statement in the model

* Point users to format_shapes utility

* Fix failing tests

* Fix remaining failing tests

* Fix further tests

* Fix a bug at AutoNormal that makes TraceMeanField test fail

Author: fehiepsi

notebooks/source/bad_posterior_geometry.ipynb

@@ -141,9 +141,13 @@
     "Instead of \n",
     "\n",
     "$$ \\beta \\sim {\\rm Normal}(0, \\lambda \\tau) $$\n",
+    "\n",
     "we write\n",
+    "\n",
     "$$ \\beta^\\prime \\sim {\\rm Normal}(0, 1) $$\n",
+    "\n",
     "and\n",
+    "\n",
     "$$ \\beta \\equiv \\lambda \\tau \\beta^\\prime  $$\n",
     "\n",
     "where $\\beta$ is now defined *deterministically* in terms of $\\lambda$, $\\tau$,\n",

numpyro/infer/autoguide.py

@@ -306,7 +306,7 @@ def __call__(self, *args, **kwargs):
                 site_fn = dist.Normal(site_loc, site_scale).to_event(event_dim)
                 if site["fn"].support is constraints.real or (
                     isinstance(site["fn"].support, constraints.independent)
-                    and site["fn"].support is constraints.real
+                    and site["fn"].support.base_constraint is constraints.real
                 ):
                     result[name] = numpyro.sample(name, site_fn)
                 else:

numpyro/util.py

@@ -510,10 +510,8 @@ def model(*args, **kwargs):
 
 def check_model_guide_match(model_trace, guide_trace):
     """
-    :param dict model_trace: The model trace to check.
-    :param dict guide_trace: The guide trace to check.
-    :raises: RuntimeWarning, ValueError
     Checks the following assumptions:
+
     1. Each sample site in the model also appears in the guide and is not
         marked auxiliary.
     2. Each sample site in the guide either appears in the model or is marked,
@@ -522,6 +520,10 @@ def check_model_guide_match(model_trace, guide_trace):
         appears in the model.
     4. At each sample site that appears in both the model and guide, the model
         and guide agree on sample shape.
+
+    :param dict model_trace: The model trace to check.
+    :param dict guide_trace: The guide trace to check.
+    :raises: RuntimeWarning, ValueError
     """
     # Check ordinary sample sites.
     guide_vars = set(
@@ -606,6 +608,28 @@ def check_model_guide_match(model_trace, guide_trace):
             )
         )
 
+    # Check if plate is missing in the model.
+    for name, site in model_trace.items():
+        if site["type"] == "sample":
+            value_ndim = jnp.ndim(site["value"])
+            batch_shape = lax.broadcast_shapes(
+                site["fn"].batch_shape,
+                jnp.shape(site["value"])[: value_ndim - len(site["fn"].event_shape)],
+            )
+            plate_dims = set(f.dim for f in site["cond_indep_stack"])
+            batch_ndim = len(batch_shape)
+            for i in range(batch_ndim):
+                dim = -i - 1
+                if batch_shape[dim] > 1 and (dim not in plate_dims):
+                    # Skip checking if it is the `scan` dimension.
+                    if dim == -batch_ndim and site.get("_control_flow_done", False):
+                        continue
+                    warnings.warn(
+                        f"Missing a plate statement for batch dimension {dim}"
+                        f" at site '{name}'. You can use `numpyro.util.format_shapes`"
+                        " utility to check shapes at all sites of your model."
+                    )
+
 
 def _format_table(rows):
     """

test/contrib/test_optim.py

@@ -69,7 +69,8 @@ def test_beta_bernoulli(elbo):
 
     def model(data):
         f = numpyro.sample("beta", dist.Beta(1.0, 1.0))
-        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Bernoulli(f), obs=data)
 
     def guide(data):
         alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive)
@@ -100,7 +101,8 @@ def test_jitted_update_fn():
 
     def model(data):
         f = numpyro.sample("beta", dist.Beta(1.0, 1.0))
-        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Bernoulli(f), obs=data)
 
     def guide(data):
         alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive)

test/infer/test_autoguide.py

@@ -69,10 +69,12 @@
 )
 def test_beta_bernoulli(auto_class):
     data = jnp.array([[1.0] * 8 + [0.0] * 2, [1.0] * 4 + [0.0] * 6]).T
+    N = len(data)
 
     def model(data):
-        f = numpyro.sample("beta", dist.Beta(jnp.ones(2), jnp.ones(2)))
-        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+        f = numpyro.sample("beta", dist.Beta(jnp.ones(2), jnp.ones(2)).to_event())
+        with numpyro.plate("N", N):
+            numpyro.sample("obs", dist.Bernoulli(f).to_event(1), obs=data)
 
     adam = optim.Adam(0.01)
     if auto_class == AutoDAIS:
@@ -104,12 +106,12 @@ def body_fn(i, val):
     # Predictive can be instantiated from posterior samples...
     predictive = Predictive(model, posterior_samples=posterior_samples)
     predictive_samples = predictive(random.PRNGKey(1), None)
-    assert predictive_samples["obs"].shape == (1000, 2)
+    assert predictive_samples["obs"].shape == (1000, N, 2)
 
     # ... or from the guide + params
     predictive = Predictive(model, guide=guide, params=params, num_samples=1000)
     predictive_samples = predictive(random.PRNGKey(1), None)
-    assert predictive_samples["obs"].shape == (1000, 2)
+    assert predictive_samples["obs"].shape == (1000, N, 2)
 
 
 @pytest.mark.parametrize(
@@ -135,9 +137,10 @@ def test_logistic_regression(auto_class, Elbo):
     labels = dist.Bernoulli(logits=logits).sample(random.PRNGKey(1))
 
     def model(data, labels):
-        coefs = numpyro.sample("coefs", dist.Normal(jnp.zeros(dim), jnp.ones(dim)))
+        coefs = numpyro.sample("coefs", dist.Normal(0, 1).expand([dim]).to_event())
         logits = numpyro.deterministic("logits", jnp.sum(coefs * data, axis=-1))
-        return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
+        with numpyro.plate("N", len(data)):
+            return numpyro.sample("obs", dist.Bernoulli(logits=logits), obs=labels)
 
     adam = optim.Adam(0.01)
     rng_key_init = random.PRNGKey(1)
@@ -242,7 +245,8 @@ def model(data):
                     dist.Uniform(0, 1), transforms.AffineTransform(0, alpha)
                 ),
             )
-        numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
 
     adam = optim.Adam(0.01)
     rng_key_init = random.PRNGKey(1)
@@ -317,12 +321,14 @@ def actual_model(data):
                     dist.Uniform(0, 1), transforms.AffineTransform(0, alpha)
                 ),
             )
-        numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
 
     def expected_model(data):
         alpha = numpyro.sample("alpha", dist.Uniform(0, 1))
         loc = numpyro.sample("loc", dist.Uniform(0, 1)) * alpha
-        numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Normal(loc, 0.1), obs=data)
 
     adam = optim.Adam(0.01)
     rng_key_init = random.PRNGKey(1)
@@ -355,9 +361,10 @@ def expected_model(data):
 def test_laplace_approximation_warning():
     def model(x, y):
         a = numpyro.sample("a", dist.Normal(0, 10))
-        b = numpyro.sample("b", dist.Normal(0, 10), sample_shape=(3,))
+        b = numpyro.sample("b", dist.Normal(0, 10).expand([3]).to_event())
         mu = a + b[0] * x + b[1] * x ** 2 + b[2] * x ** 3
-        numpyro.sample("y", dist.Normal(mu, 0.001), obs=y)
+        with numpyro.plate("N", len(x)):
+            numpyro.sample("y", dist.Normal(mu, 0.001), obs=y)
 
     x = random.normal(random.PRNGKey(0), (3,))
     y = 1 + 2 * x + 3 * x ** 2 + 4 * x ** 3
@@ -375,7 +382,8 @@ def model(x, y):
         a = numpyro.sample("a", dist.Normal(0, 10))
         b = numpyro.sample("b", dist.Normal(0, 10))
         mu = a + b * x
-        numpyro.sample("y", dist.Normal(mu, 1), obs=y)
+        with numpyro.plate("N", len(x)):
+            numpyro.sample("y", dist.Normal(mu, 1), obs=y)
 
     x = random.normal(random.PRNGKey(0), (100,))
     y = 1 + 2 * x
@@ -401,7 +409,8 @@ def model(y):
             "sigma", dist.ImproperUniform(dist.constraints.positive, (), ())
         )
         mu = numpyro.deterministic("mu", lambda1 + lambda2)
-        numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
+        with numpyro.plate("N", len(y)):
+            numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
 
     guide = AutoDiagonalNormal(model)
     svi = SVI(model, guide, optim.Adam(0.003), Trace_ELBO(), y=y)
@@ -417,7 +426,8 @@ def model(x, y):
         nn = numpyro.module("nn", Dense(1), (10,))
         mu = nn(x).squeeze(-1)
         sigma = numpyro.sample("sigma", dist.HalfNormal(1))
-        numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
+        with numpyro.plate("N", len(y)):
+            numpyro.sample("y", dist.Normal(mu, sigma), obs=y)
 
     guide = AutoDiagonalNormal(model)
     svi = SVI(model, guide, optim.Adam(0.003), Trace_ELBO(), x=x, y=y)
@@ -497,7 +507,8 @@ def model(y=None):
         mu = numpyro.sample("mu", dist.Normal(0, 5))
         sigma = numpyro.param("sigma", 1, constraint=constraints.positive)
 
-        y = numpyro.sample("y", dist.Normal(mu, sigma).expand((n,)), obs=y)
+        with numpyro.plate("N", len(y)):
+            y = numpyro.sample("y", dist.Normal(mu, sigma).expand((n,)), obs=y)
         numpyro.deterministic("z", (y - mu) / sigma)
 
     mu, sigma = 2, 3

test/infer/test_svi.py

@@ -63,7 +63,8 @@ def test_beta_bernoulli(elbo, optimizer):
 
     def model(data):
         f = numpyro.sample("beta", dist.Beta(1.0, 1.0))
-        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Bernoulli(f), obs=data)
 
     def guide(data):
         alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive)
@@ -94,7 +95,8 @@ def test_run(progress_bar):
 
     def model(data):
         f = numpyro.sample("beta", dist.Beta(1.0, 1.0))
-        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Bernoulli(f), obs=data)
 
     def guide(data):
         alpha_q = numpyro.param(
@@ -124,7 +126,8 @@ def test_jitted_update_fn():
 
     def model(data):
         f = numpyro.sample("beta", dist.Beta(1.0, 1.0))
-        numpyro.sample("obs", dist.Bernoulli(f), obs=data)
+        with numpyro.plate("N", len(data)):
+            numpyro.sample("obs", dist.Bernoulli(f), obs=data)
 
     def guide(data):
         alpha_q = numpyro.param("alpha_q", 1.0, constraint=constraints.positive)

test/test_util.py

@@ -170,28 +170,27 @@ def model_test():
     )
 
 
-def test_check_model_guide_match():
-    def _run_svi(model, guide):
+def _run_svi_check_warnings(model, guide, expected_string):
+    with pytest.warns(UserWarning, match=expected_string) as ws:
         adam = numpyro.optim.Adam(1e-3)
         svi = numpyro.infer.SVI(model, guide, adam, numpyro.infer.Trace_ELBO())
-        svi.run(random.PRNGKey(42), num_steps=50)
-
-    def _run_svi_check_warnings(model, guide, expected_string):
-        with pytest.warns(UserWarning, match=expected_string) as ws:
-            _run_svi(model, guide)
-            assert len(ws) == 1
-            assert expected_string in str(ws[0].message)
-
-    def _create_traces_check_error_string(model, guide, expected_string):
-        model_trace = numpyro.handlers.trace(
-            numpyro.handlers.seed(model, rng_seed=42)
-        ).get_trace()
-        guide_trace = numpyro.handlers.trace(
-            numpyro.handlers.seed(guide, rng_seed=42)
-        ).get_trace()
-        with pytest.raises(ValueError, match=expected_string):
-            check_model_guide_match(model_trace, guide_trace)
+        svi.run(random.PRNGKey(42), num_steps=5)
+        assert len(ws) == 1
+        assert expected_string in str(ws[0].message)
+
+
+def _create_traces_check_error_string(model, guide, expected_string):
+    model_trace = numpyro.handlers.trace(
+        numpyro.handlers.seed(model, rng_seed=42)
+    ).get_trace()
+    guide_trace = numpyro.handlers.trace(
+        numpyro.handlers.seed(guide, rng_seed=42)
+    ).get_trace()
+    with pytest.raises(ValueError, match=expected_string):
+        check_model_guide_match(model_trace, guide_trace)
 
+
+def test_check_model_guide_match():
     # 1. Auxiliary vars in the model
     def model():
         numpyro.sample("x", dist.Normal())
@@ -236,7 +235,9 @@ def guide():
 
     # 5. Check shapes agree
     def model():
-        numpyro.sample("x", dist.Normal().expand((3, 2)))
+        with numpyro.plate("a", 3, dim=-2):
+            with numpyro.plate("b", 2, dim=-1):
+                numpyro.sample("x", dist.Normal().expand((3, 2)))
 
     def guide():
         numpyro.sample("x", dist.Normal().expand((3, 5)))
@@ -245,12 +246,24 @@ def guide():
 
     # 6. Check subsample sites introduced by plate
     def model():
-        numpyro.sample("x", dist.Normal().expand((10,)))
+        with numpyro.plate("a", 10):
+            numpyro.sample("x", dist.Normal().expand((10,)))
 
     def guide():
-        with numpyro.handlers.plate("data", 100, subsample_size=10):
+        with numpyro.plate("data", 100, subsample_size=10):
             numpyro.sample("x", dist.Normal())
 
     _run_svi_check_warnings(
         model, guide, "Found plate statements in guide but not model"
     )
+
+
+def test_missing_plate_in_model():
+    def model():
+        x = numpyro.sample("x", dist.Normal(0, 1))
+        numpyro.sample("obs", dist.Normal(x, 1), obs=jnp.ones(10))
+
+    def guide():
+        numpyro.sample("x", dist.Normal(0, 1))
+
+    _run_svi_check_warnings(model, guide, "Missing a plate statement")