Repository: pyro-ppl/numpyro

Issue:
A way to notify shapes going wrong in log_prob computation in ELBO?
We've recently run into a subtle (modelling) bug with broadcasting in the ELBO loss.
The offending part in our model boiled down to something like this:
```
import numpyro.distributions as dists

from numpyro import plate, sample
from numpyro.handlers import trace, seed

import numpy as np

N = 100
d = 7
logits = np.random.randn(N)
ys = np.random.binomial(1, .5, size=(N, 1))
 
def model(logits, ys):
    with plate("batch", len(ys)):
        sample("ys", dists.Bernoulli(logits), obs=ys)

# [some SVI inference code with Trace_ELBO here]
 ```
 
Note that while `logits` are of shape `(N,)`, our `ys` had shape `(N,1)` from our preprocessing script. As a consequence, the ELBO-internal call to `Bernoulli(logits).log_prob(ys)`  broadcasts over `ys` and returns a `(N, N)` matrix, over which the ELBO / `log_density` then sums. The program does not produce any error but our inference obviously fails.

Now, this is an oversight by us and not strictly a bug in numpyro (there could be scenarios where this behavior could be desired by a user, I guess (although I find it unlikely)). However, due the log_prob being abstracted away so nicely, it is very difficult to track this down when it is not desired, as in our case (we have been struggling to solve our poor model performance for a long time before we actually figured out that this is happening). I therefore wanted to bring this up and ideally have a discussion on whether there is a way to make this easier to spot if it happens.

One thought I had off the top of my head was to add an optional `expected_log_prob_shape` kwargs to `sample` where the user can specify what shape they would expect coming out of `log_prob` and the program failing if this is not the case. There might even be a way to do this without the explicit argument based on the distributions batch and even shapes maybe? I think another useful thing that could help a bit would be to actually check given `obs` for shape consistency with `sample_shape` (if both are provided to `sample`) as a sanity check for the user.
TraceGraph_ELBO bug shape error
Hi, I've been running into a bug related to data with multiple likelihoods. It seems that when computing downstream costs in the TraceGraph, multiple children have an empty `frozenset` frames on line 389 of `elbo.py` resulting in shape mismatch due to lookup bugs in the `MultiFrameTensor`.

I have an example program [here](https://gist.github.com/quattro/c3cdf41bf52a2334aa2dd2851388d254) to recreate the issue. The example I provided is a trivial example where the data likelihoods are completely independent (and thus could be optimized separately), but this error also occurs in the situation when the data likelihoods are conditionally independent given some latent variable.

Ideally these types of situations would have another dimension to reflect different data settings with similarly structured likelihoods (and possible 'resolve' this issue naturally), but the shapes change per data type and don't work cleanly in a single tensor structure.

```
Traceback (most recent call last):
  File "/Users/nicholas/Projects/VAE_SuSiE/g_test.py", line 180, in <module>
    sys.exit(main(sys.argv[1:]))
  File "/Users/nicholas/Projects/VAE_SuSiE/g_test.py", line 166, in main
    results = svi.run(
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/svi.py", line 327, in run
    svi_state, loss = jit(body_fn)(svi_state, None)
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/svi.py", line 316, in body_fn
    svi_state, loss = self.stable_update(svi_state, *args, **kwargs)
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/svi.py", line 276, in stable_update
    (loss_val, mutable_state), optim_state = self.optim.eval_and_stable_update(
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/optim.py", line 93, in eval_and_stable_update
    (out, aux), grads = value_and_grad(fn, has_aux=True)(params)
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/svi.py", line 59, in loss_fn
    elbo.loss(
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/elbo.py", line 604, in loss
    return -single_particle_elbo(rng_key)
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/elbo.py", line 579, in single_particle_elbo
    downstream_costs, _ = _compute_downstream_costs(
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/elbo.py", line 483, in _compute_downstream_costs
    downstream_costs[node].add(*downstream_costs[child].items())
  File "/Users/nicholas/miniforge3/lib/python3.9/site-packages/numpyro/infer/elbo.py", line 389, in add
    self[frames] = self[frames] + value
  File "/Users/nicholas/Projects/jax/jax/_src/numpy/lax_numpy.py", line 5847, in deferring_binary_op
    return binary_op(self, other)
  File "/Users/nicholas/Projects/jax/jax/_src/numpy/lax_numpy.py", line 431, in fn
    return lax_fn(x1, x2) if x1.dtype != bool_ else bool_lax_fn(x1, x2)
TypeError: add got incompatible shapes for broadcasting: (2,), (4,).
```


Discussion:
i'd suggest using the [format_shapes](http://num.pyro.ai/en/latest/utilities.html?highlight=%20format_shapes#numpyro.util.format_shapes) utility function. see [here](https://github.com/pyro-ppl/numpyro/blob/145ed8373fe268798b352b98350ffc596a51e682/test/test_util.py#L149) for example usage.

i guess one possibility would be to have an argument to e.g. `SVI` that prints this information to stdout automatically once upon initialization
I think @maw501 already addresses part of the shape-mismatched issue in #1171. One way is to enhance the utility that @maw501 added and raise a warning whenever user code has missing plate statements like in the case
```
    with plate("batch", len(ys)):
        sample("ys", dists.Bernoulli(logits), obs=ys)
```
where `ys` has rank 2 but only one plate is declared.
> have an argument to e.g. SVI that prints this information to stdout automatically once upon initialization

Good idea! I guess when a plate is missing, we can print out this message together with the warning. (When all plates are correctly declared, I guess `format_shapes` is not useful)
Hi @quattro, could you add `to_event` to the latent sites such as `Normal` (or put them under some `plate`)? I guess that causes the issue.
Hi @fehiepsi , apologies. While `to_event(1)` didn't seem to correct the issue, wrapping `beta~Normal` in a `plate` with `size=1` got it to run.

Would it be possible to catch and re-throw a more informative message in this setting (ie conditional independence is unclear)?
That's a good point. I think we should throw errors when `plate` is missing. Looking at the code base, it is not clear to me where we should add the check. I'll follow up later. :)

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/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:
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/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 patch:
diff --git a/test/contrib/test_optim.py b/test/contrib/test_optim.py
--- a/test/contrib/test_optim.py
+++ b/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)
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
@@ -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
diff --git a/test/infer/test_svi.py b/test/infer/test_svi.py
--- a/test/infer/test_svi.py
+++ b/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)
diff --git a/test/test_util.py b/test/test_util.py
--- a/test/test_util.py
+++ b/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")