Fix sample_prior_predictive edge case (#3048)

Author: ColCarroll

pymc3/distributions/distribution.py

@@ -470,7 +470,14 @@ def generate_samples(generator, *args, **kwargs):
 
     if broadcast_shape is None:
         inputs = args + tuple(kwargs.values())
-        broadcast_shape = np.broadcast(*inputs).shape  # size of generator(size=1)
+        try:
+            broadcast_shape = np.broadcast(*inputs).shape  # size of generator(size=1)
+        except ValueError:  # sometimes happens if args have shape (500,) and (500, 4)
+            max_dims = max(j.ndim for j in args + tuple(kwargs.values()))
+            args = tuple([j.reshape(j.shape + (1,) * (max_dims - j.ndim)) for j in args])
+            kwargs = {k: v.reshape(v.shape + (1,) * (max_dims - v.ndim)) for k, v in kwargs.items()}
+            inputs = args + tuple(kwargs.values())
+            broadcast_shape = np.broadcast(*inputs).shape  # size of generator(size=1)
 
     dist_shape = to_tuple(dist_shape)
     broadcast_shape = to_tuple(broadcast_shape)

pymc3/tests/test_sampling.py

@@ -432,3 +432,11 @@ def test_density_dist(self):
             prior = pm.sample_prior_predictive()
 
         npt.assert_almost_equal(prior['a'].mean(), 0, decimal=1)
+
+    def test_shape_edgecase(self):
+        with pm.Model():
+            mu = pm.Normal('mu', shape=5)
+            sd = pm.Uniform('sd', lower=2, upper=3)
+            x = pm.Normal('x', mu=mu, sd=sd, shape=5)
+            prior = pm.sample_prior_predictive(10)
+        assert prior['mu'].shape == (10, 5)