Check plate consistent in auto guide (#1049)

* check plate consistent and add regression test

* make format

Author: fehiepsi

numpyro/infer/autoguide.py

@@ -148,7 +148,12 @@ def _setup_prototype(self, *args, **kwargs):
         for name, site in self.prototype_trace.items():
             if site["type"] == "sample":
                 for frame in site["cond_indep_stack"]:
-                    self._prototype_frames[frame.name] = frame
+                    if frame.name in self._prototype_frames:
+                        assert (
+                            frame == self._prototype_frames[frame.name]
+                        ), f"The plate {frame.name} has inconsistent dim or size. Please check your model again."
+                    else:
+                        self._prototype_frames[frame.name] = frame
             elif site["type"] == "plate":
                 self._prototype_frame_full_sizes[name] = site["args"][0]
 

test/infer/test_autoguide.py

@@ -482,3 +482,17 @@ def model(y=None):
         predictive_samples["z"],
         atol=0.05,
     )
+
+
+@pytest.mark.parametrize("size,dim", [(10, -2), (5, -1)])
+def test_plate_inconsistent(size, dim):
+    def model():
+        with numpyro.plate("a", 10, dim=-1):
+            numpyro.sample("x", dist.Normal(0, 1))
+        with numpyro.plate("a", size, dim=dim):
+            numpyro.sample("y", dist.Normal(0, 1))
+
+    guide = AutoDelta(model)
+    svi = SVI(model, guide, numpyro.optim.Adam(step_size=0.1), Trace_ELBO())
+    with pytest.raises(AssertionError, match="has inconsistent dim or size"):
+        svi.run(random.PRNGKey(0), 10)