Allow arbitrary order of plate (#555)

Author: fehiepsi

numpyro/primitives.py

@@ -265,10 +265,7 @@ def process_message(self, msg):
         cond_indep_stack.append(frame)
         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)))
+        overlap_idx = max(len(expected_shape) - len(dist_batch_shape), 0)
         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)

test/test_handlers.py

@@ -166,6 +166,16 @@ def model_nested_plates_2():
         assert xy.shape == (5, 1, 10)
 
 
+def model_nested_plates_3():
+    outer = numpyro.plate('outer', 10, dim=-1)
+    inner = numpyro.plate('inner', 5, dim=-2)
+    numpyro.deterministic('z', 1.)
+
+    with inner, outer:
+        xy = numpyro.sample('xy', dist.Normal(np.zeros((5, 10)), 1.))
+        assert xy.shape == (5, 10)
+
+
 def model_dist_batch_shape():
     outer = numpyro.plate('outer', 10)
     inner = numpyro.plate('inner', 5, dim=-3)
@@ -204,6 +214,7 @@ def model_subsample_1():
     model_nested_plates_0,
     model_nested_plates_1,
     model_nested_plates_2,
+    model_nested_plates_3,
     model_dist_batch_shape,
     model_subsample_1,
 ])