Repository: pyro-ppl/numpyro

Issue:
`plate` has unexpected behaviour when dist has batch_shape
I faced this issue while revising baseball example using `plate`. The repro code
```
import jax.numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.handlers import trace, seed

def model():
    with numpyro.plate('plate', 10):
        numpyro.sample('x', dist.Normal(np.zeros(10), 1))

trace(seed(model, 0)).get_trace()['x']['value'].shape
```
should return `(10,)` but currently, it returns `(10, 10)`.

cc @neerajprad 


Discussion:
Thanks for the bug report! We need to handle the interaction between batch and sample shapes better. I'll fix this shortly. 
Hmm..this issue is actually more complicated. Currently we can only change a distribution's shape on the left. If we have a distribution with batch shape `[2, 1, 10]`, we cannot just expand it to some arbitrary batch shape like `[2, 5, 10]`, unlike in Pyro which allows for this using `.expand`. The case above is more easily handled though. I'll address that first.

Resolving patch:
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -231,7 +231,20 @@ def process_message(self, msg):
         cond_indep_stack = msg['cond_indep_stack']
         frame = CondIndepStackFrame(self.name, self.dim, self.subsample_size)
         cond_indep_stack.append(frame)
-        batch_shape = self._get_batch_shape(cond_indep_stack)
+        expected_shape = self._get_batch_shape(cond_indep_stack)
+        dist_batch_shape = msg['fn'].batch_shape if msg['type'] == 'sample' else ()
+        overlap_idx = len(expected_shape) - len(dist_batch_shape)
+        if overlap_idx < 0:
+            raise ValueError('Expected dimensions within plate = {}, which is less than the '
+                             'distribution\'s batch shape = {}.'.format(len(expected_shape), len(dist_batch_shape)))
+        trailing_shape = expected_shape[overlap_idx:]
+        # e.g. distribution with batch shape (1, 5) cannot be broadcast to (5, 5)
+        broadcast_shape = lax.broadcast_shapes(trailing_shape, dist_batch_shape)
+        if broadcast_shape != dist_batch_shape:
+            raise ValueError('Distribution batch shape = {} cannot be broadcast up to {}. '
+                             'Consider using unbatched distributions.'
+                             .format(dist_batch_shape, broadcast_shape))
+        batch_shape = expected_shape[:overlap_idx]
         if 'sample_shape' in msg['kwargs']:
             batch_shape = lax.broadcast_shapes(msg['kwargs']['sample_shape'], batch_shape)
         msg['kwargs']['sample_shape'] = batch_shape


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -110,6 +110,21 @@ def model_nested_plates_2():
         assert xy.shape == (5, 1, 10)
 
 
+def model_dist_batch_shape():
+    outer = numpyro.plate('outer', 10)
+    inner = numpyro.plate('inner', 5, dim=-3)
+    with outer:
+        x = numpyro.sample('x', dist.Normal(np.zeros(10), 1.))
+        assert x.shape == (10,)
+    with inner:
+        y = numpyro.sample('y', dist.Normal(0., np.ones(10)))
+        assert y.shape == (5, 1, 10)
+
+    with outer, inner:
+        xy = numpyro.sample('xy', dist.Normal(0., np.ones(10)), sample_shape=(10,))
+        assert xy.shape == (5, 10, 10)
+
+
 def model_subsample_1():
     outer = numpyro.plate('outer', 20, subsample_size=10)
     inner = numpyro.plate('inner', 10, subsample_size=5, dim=-3)
@@ -129,6 +144,7 @@ def model_subsample_1():
     model_nested_plates_0,
     model_nested_plates_1,
     model_nested_plates_2,
+    model_dist_batch_shape,
     model_subsample_1,
 ])
 def test_plate(model):
diff --git a/test/test_infer_util.py b/test/test_infer_util.py
--- a/test/test_infer_util.py
+++ b/test/test_infer_util.py
@@ -17,7 +17,7 @@ def beta_bernoulli():
 
     def model(data=None):
         beta = numpyro.sample("beta", dist.Beta(np.ones(5), np.ones(5)))
-        with numpyro.plate("plate", 1000):
+        with numpyro.plate("plate", 1000, dim=-2):
             numpyro.sample("obs", dist.Bernoulli(beta), obs=data)
 
     return model, data, true_probs