Repository: pyro-ppl/numpyro

Issue:
FR: Fail with informative error message when same plate name used with inconsistent dimensions
Currently reusing the same plate name with different arguments to e.g. dim will cause problems. Models like this could be rejected with an informative error message quickly, before confusing problems which can include arrays changing shape happen. See below for the current behaviour. 

```py
import sys

import numpyro
import numpyro.distributions as dist
from jax import numpy as jnp
from jax import random
from numpyro.infer import SVI, Trace_ELBO
from numpyro.infer.autoguide import AutoDelta
from numpyro.infer.initialization import init_to_median
from numpyro.optim import Adam


def irt2pl(ncls, resp, word_subsample_size=None):
    nstud, nitems = resp.shape
    difficulty_offsets = numpyro.sample(
        "difficulty_offsets",
        dist.TransformedDistribution(
            dist.Normal(0, 1).expand([ncls - 1]), dist.transforms.OrderedTransform()
        ),
    )
    with numpyro.plate("nstud", nstud, dim=-1):
        abilities = numpyro.sample("abilities", dist.Normal())
    with numpyro.plate("nitems", nitems, dim=-1):
        difficulties = numpyro.sample("difficulties", dist.Normal())
        discriminations = numpyro.sample("discriminations", dist.HalfNormal())
    offset_difficulties = jnp.expand_dims(difficulties, 1) + jnp.expand_dims(
        difficulty_offsets, 0
    )
    print("abilities.shape", abilities.shape)
    print("discriminations.shape", discriminations.shape)
    predictor = jnp.expand_dims(abilities, 1) * jnp.expand_dims(discriminations, 0)
    cutpoints = jnp.expand_dims(
        offset_difficulties * jnp.expand_dims(discriminations, 0), 0
    )
    with numpyro.plate("nstud", nstud, dim=-2), numpyro.plate("nitems", nitems, dim=-1):
        numpyro.sample("resp", dist.OrderedLogistic(predictor, cutpoints), obs=resp)


resp = jnp.array(
    [[0, 1, 4, 3], [0, 1, 4, 3], [3, 4, 4, 4], [2, 2, 4, 4], [2, 2, 4, 4], [1, 2, 3, 3]]
)

optim = Adam(0.1, 0.8, 0.99)
elbo = Trace_ELBO()
guide = AutoDelta(irt2pl, init_loc_fn=init_to_median())
rng_key = random.PRNGKey(42)
svi = SVI(irt2pl, guide, optim, loss=elbo)
svi.run(rng_key, 200, 5, resp)
print(guide(5, resp))
```

Running gets an error like:
```
abilities.shape (6,)
discriminations.shape (4,)
abilities.shape (6,)
discriminations.shape (4,)
abilities.shape (6,)
discriminations.shape (4,)
abilities.shape (6,)
discriminations.shape (4,)
abilities.shape (6,)
discriminations.shape (4,)
abilities.shape (6, 6)
discriminations.shape (4,)
Traceback (most recent call last):
  File "broken.py", line 55, in <module>
    svi.run(rng_key, 200, 5, resp)
...
  File "broken.py", line 31, in irt2pl
    predictor = jnp.expand_dims(abilities, 1) * jnp.expand_dims(discriminations, 0)
  File "/path/to/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 5256, in deferring_binary_op
    return binary_op(self, other)
  File "/home/frankier/edu/doc/vocabmodel/.direnv/python-3.8.6/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 396, in fn
    x1, x2 = _promote_args(numpy_fn.__name__, x1, x2)
  File "/path/to/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 333, in _promote_args
    return _promote_shapes(fun_name, *_promote_dtypes(*args))
  File "/path/to/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 251, in _promote_shapes
    result_rank = len(lax.broadcast_shapes(*shapes))
  File "/path/to/python3.8/site-packages/jax/_src/util.py", line 198, in wrapper
    return cached(bool(config.x64_enabled), *args, **kwargs)
  File "/path/to/python3.8/site-packages/jax/_src/util.py", line 191, in cached
    return f(*args, **kwargs)
  File "/path/to/python3.8/site-packages/jax/_src/lax/lax.py", line 97, in broadcast_shapes
    raise ValueError("Incompatible shapes for broadcasting: {}"
ValueError: Incompatible shapes for broadcasting: ((6, 1, 6), (1, 1, 4))
```


Discussion:

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/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 patch:
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
@@ -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)