Repository: pymc-devs/pymc
Pull request: Fix for #3310 and some more
Added `broadcast_distribution_samples`, which helps broadcasting multiple rvs calls with different `size` and distribution parameter shapes. Added shape guards to other continuous distributions. In particular, this fixes #3310, but it also prevents similar things from happening in other discrete and continuous distributions that could encounter the same pathological behavior.

Problem statement:
Shape issue when prior predictive sampling from ZeroInflatedPoisson
```python
with pm.Model() as model:
    θ = pm.Beta('θ', alpha=1, beta=1)
    ψ = pm.HalfNormal('ψ', sd=1)
    s = pm.ZeroInflatedPoisson('suppliers', psi=ψ, theta=θ, 
                               shape=20)
    gen_data = pm.sample_prior_predictive(samples=5000)
```

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-24-aef3a9b7b63a> in <module>()
      4     s = pm.ZeroInflatedPoisson('suppliers', psi=ψ, theta=θ, 
      5                                shape=20)
----> 6     gen_data = pm.sample_prior_predictive(samples=5000)

~/working/projects/pymc/pymc3/sampling.py in sample_prior_predictive(samples, model, vars, random_seed)
   1323     names = get_default_varnames(model.named_vars, include_transformed=False)
   1324     # draw_values fails with auto-transformed variables. transform them later!
-> 1325     values = draw_values([model[name] for name in names], size=samples)
   1326 
   1327     data = {k: v for k, v in zip(names, values)}

~/working/projects/pymc/pymc3/distributions/distribution.py in draw_values(params, point, size)
    374                     # the stack of nodes to try to draw from. We exclude the
    375                     # nodes in the `params` list.
--> 376                     stack.extend([node for node in named_nodes_parents[next_]
    377                                   if node is not None and
    378                                   node.name not in drawn and

~/working/projects/pymc/pymc3/distributions/distribution.py in _draw_value(param, point, givens, size)
    468                 # shape inspection for ObservedRV
    469                 dist_tmp = param.distribution
--> 470                 try:
    471                     distshape = param.observations.shape.eval()
    472                 except AttributeError:

~/working/projects/pymc/pymc3/model.py in __call__(self, *args, **kwargs)
     41 
     42     def __call__(self, *args, **kwargs):
---> 43         return getattr(self.obj, self.method_name)(*args, **kwargs)
     44 
     45 

~/working/projects/pymc/pymc3/distributions/discrete.py in random(self, point, size)
    848                              dist_shape=self.shape,
    849                              size=size)
--> 850         return g * (np.random.random(np.squeeze(g.shape)) < psi)
    851 
    852     def logp(self, value):

ValueError: operands could not be broadcast together with shapes (5000,20) (5000,) 
```

Hints from issue discussion:
CC @lucianopaz 
theta and psi are scalars but when we draw 5000 samples from their distribution, we get them as 1D arrays of shape `(5000,)` instead of `(5000, 1)` which would broadcast well. I'll check to see what we need to do to get these things to work
This is kind of a recurrent problem that we face with distributions that inside of their `random` method do many calls to different random number generators apart from the main `generate_samples`. All `ZeroInflated*` classes have the same problem. The `Mixture` class has a similar but harder problem (#3270). The `Multinomial` class had a similar problem (#3271). The thing is that `generate_samples` does some logic on how things should broadcast and returns samples of the correct shape in the end, but the `random` method should also rebroadcast the distribution parameters being aware of the `size` input. I'll try to make a more generic fix than just patching up `ZeroInflatedPoisson`.

Commit message: Fixed lint error

Patch to review:
diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md
index 8a6fdbf358..7df4a49981 100644
--- a/RELEASE-NOTES.md
+++ b/RELEASE-NOTES.md
@@ -36,6 +36,8 @@ This will be the last release to support Python 2.
 - Fixed `Rice` distribution, which inconsistently mixed two parametrizations (#3286).
 - `Rice` distribution now accepts multiple parameters and observations and is usable with NUTS (#3289).
 - `sample_posterior_predictive` no longer calls `draw_values` to initialize the shape of the ppc trace. This called could lead to `ValueError`'s when sampling the ppc from a model with `Flat` or `HalfFlat` prior distributions (Fix issue #3294).
+- Added the `broadcast_distribution_samples` function that helps broadcasting arrays of drawn samples, taking into account the requested `size` and the inferred distribution shape. This sometimes is needed by distributions that call several `rvs` separately within their `random` method, such as the `ZeroInflatedPoisson` (Fix issue #3310).
+- The `Wald`, `Kumaraswamy`, `LogNormal`, `Pareto`, `Cauchy`, `HalfCauchy`, `Weibull` and `ExGaussian` distributions `random` method used a hidden `_random` function that was written with scalars in mind. This could potentially lead to artificial correlations between random draws. Added shape guards and broadcasting of the distribution samples to prevent this (Similar to issue #3310).
 
 
 ### Deprecations
diff --git a/pymc3/distributions/continuous.py b/pymc3/distributions/continuous.py
index 91c52e38c4..9e214a4bff 100644
--- a/pymc3/distributions/continuous.py
+++ b/pymc3/distributions/continuous.py
@@ -21,7 +21,8 @@
     alltrue_elemwise, betaln, bound, gammaln, i0e, incomplete_beta, logpow,
     normal_lccdf, normal_lcdf, SplineWrapper, std_cdf, zvalue,
 )
-from .distribution import Continuous, draw_values, generate_samples
+from .distribution import (Continuous, draw_values, generate_samples,
+                           broadcast_distribution_samples)
 
 __all__ = ['Uniform', 'Flat', 'HalfFlat', 'Normal', 'TruncatedNormal', 'Beta',
            'Kumaraswamy', 'Exponential', 'Laplace', 'StudentT', 'Cauchy',
@@ -964,6 +965,11 @@ def random(self, point=None, size=None):
         """
         mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha],
                                      point=point, size=size)
+        print(mu.shape, lam.shape, alpha.shape)
+        mu, lam, alpha = broadcast_distribution_samples([mu, lam, alpha],
+                                                        size=size)
+        print(mu.shape, lam.shape, alpha.shape)
+        print(self.shape, size)
         return generate_samples(self._random,
                                 mu, lam, alpha,
                                 dist_shape=self.shape,
@@ -1292,6 +1298,7 @@ def random(self, point=None, size=None):
         """
         a, b = draw_values([self.a, self.b],
                            point=point, size=size)
+        a, b = broadcast_distribution_samples([a, b], size=size)
         return generate_samples(self._random, a, b,
                                 dist_shape=self.shape,
                                 size=size)
@@ -1665,6 +1672,7 @@ def random(self, point=None, size=None):
         array
         """
         mu, tau = draw_values([self.mu, self.tau], point=point, size=size)
+        mu, tau = broadcast_distribution_samples([mu, tau], size=size)
         return generate_samples(self._random, mu, tau,
                                 dist_shape=self.shape,
                                 size=size)
@@ -1952,6 +1960,7 @@ def random(self, point=None, size=None):
         """
         alpha, m = draw_values([self.alpha, self.m],
                                point=point, size=size)
+        alpha, m = broadcast_distribution_samples([alpha, m], size=size)
         return generate_samples(self._random, alpha, m,
                                 dist_shape=self.shape,
                                 size=size)
@@ -2076,6 +2085,7 @@ def random(self, point=None, size=None):
         """
         alpha, beta = draw_values([self.alpha, self.beta],
                                   point=point, size=size)
+        alpha, beta = broadcast_distribution_samples([alpha, beta], size=size)
         return generate_samples(self._random, alpha, beta,
                                 dist_shape=self.shape,
                                 size=size)
@@ -2637,6 +2647,7 @@ def random(self, point=None, size=None):
         """
         alpha, beta = draw_values([self.alpha, self.beta],
                                   point=point, size=size)
+        alpha, beta = broadcast_distribution_samples([alpha, beta], size=size)
 
         def _random(a, b, size=None):
             return b * (-np.log(np.random.uniform(size=size)))**(1 / a)
@@ -2921,6 +2932,8 @@ def random(self, point=None, size=None):
         """
         mu, sigma, nu = draw_values([self.mu, self.sigma, self.nu],
                                     point=point, size=size)
+        mu, sigma, nu = broadcast_distribution_samples([mu, sigma, nu],
+                                                       size=size)
 
         def _random(mu, sigma, nu, size=None):
             return (np.random.normal(mu, sigma, size=size)
diff --git a/pymc3/distributions/discrete.py b/pymc3/distributions/discrete.py
index 86db441ae1..a959656cc1 100644
--- a/pymc3/distributions/discrete.py
+++ b/pymc3/distributions/discrete.py
@@ -6,7 +6,8 @@
 
 from pymc3.util import get_variable_name
 from .dist_math import bound, factln, binomln, betaln, logpow, random_choice
-from .distribution import Discrete, draw_values, generate_samples
+from .distribution import (Discrete, draw_values, generate_samples,
+                           broadcast_distribution_samples)
 from pymc3.math import tround, sigmoid, logaddexp, logit, log1pexp
 
 
@@ -345,6 +346,7 @@ def _ppf(self, p):
 
     def _random(self, q, beta, size=None):
         p = np.random.uniform(size=size)
+        p, q, beta = broadcast_distribution_samples([p, q, beta], size=size)
 
         return np.ceil(np.power(np.log(1 - p) / np.log(q), 1. / beta)) - 1
 
@@ -847,7 +849,8 @@ def random(self, point=None, size=None):
         g = generate_samples(stats.poisson.rvs, theta,
                              dist_shape=self.shape,
                              size=size)
-        return g * (np.random.random(np.squeeze(g.shape)) < psi)
+        g, psi = broadcast_distribution_samples([g, psi], size=size)
+        return g * (np.random.random(g.shape) < psi)
 
     def logp(self, value):
         psi = self.psi
@@ -939,7 +942,8 @@ def random(self, point=None, size=None):
         g = generate_samples(stats.binom.rvs, n, p,
                              dist_shape=self.shape,
                              size=size)
-        return g * (np.random.random(np.squeeze(g.shape)) < psi)
+        g, psi = broadcast_distribution_samples([g, psi], size=size)
+        return g * (np.random.random(g.shape) < psi)
 
     def logp(self, value):
         psi = self.psi
@@ -1057,7 +1061,8 @@ def random(self, point=None, size=None):
                              dist_shape=self.shape,
                              size=size)
         g[g == 0] = np.finfo(float).eps  # Just in case
-        return stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
+        g, psi = broadcast_distribution_samples([g, psi], size=size)
+        return stats.poisson.rvs(g) * (np.random.random(g.shape) < psi)
 
     def logp(self, value):
         alpha = self.alpha
diff --git a/pymc3/distributions/distribution.py b/pymc3/distributions/distribution.py
index 0dcb6bdb48..22ab6656fe 100644
--- a/pymc3/distributions/distribution.py
+++ b/pymc3/distributions/distribution.py
@@ -640,3 +640,30 @@ def generate_samples(generator, *args, **kwargs):
     if one_d and samples.shape[-1] == 1:
         samples = samples.reshape(samples.shape[:-1])
     return np.asarray(samples)
+
+
+def broadcast_distribution_samples(samples, size=None):
+    if size is None:
+        return np.broadcast_arrays(*samples)
+    _size = to_tuple(size)
+    try:
+        broadcasted_samples = np.broadcast_arrays(*samples)
+    except ValueError:
+        # Raw samples shapes
+        p_shapes = [p.shape for p in samples]
+        # samples shapes without the size prepend
+        sp_shapes = [s[len(_size):] if _size == s[:len(_size)] else s
+                     for s in p_shapes]
+        broadcast_shape = np.broadcast(*[np.empty(s) for s in sp_shapes]).shape
+        broadcasted_samples = []
+        for param, p_shape, sp_shape in zip(samples, p_shapes, sp_shapes):
+            if _size == p_shape[:len(_size)]:
+                slicer_head = [slice(None)] * len(_size)
+            else:
+                slicer_head = [np.newaxis] * len(_size)
+            slicer_tail = ([np.newaxis] * (len(broadcast_shape) -
+                                           len(sp_shape)) +
+                           [slice(None)] * len(sp_shape))
+            broadcasted_samples.append(param[tuple(slicer_head + slicer_tail)])
+        broadcasted_samples = np.broadcast_arrays(*broadcasted_samples)
+    return broadcasted_samples
diff --git a/pymc3/tests/test_sampling.py b/pymc3/tests/test_sampling.py
index 4b43bedad1..b626e1abaa 100644
--- a/pymc3/tests/test_sampling.py
+++ b/pymc3/tests/test_sampling.py
@@ -467,3 +467,13 @@ def test_shape_edgecase(self):
             x = pm.Normal('x', mu=mu, sd=sd, shape=5)
             prior = pm.sample_prior_predictive(10)
         assert prior['mu'].shape == (10, 5)
+
+    def test_zeroinflatedpoisson(self):
+        with pm.Model():
+            theta = pm.Beta('theta', alpha=1, beta=1)
+            psi = pm.HalfNormal('psi', sd=1)
+            pm.ZeroInflatedPoisson('suppliers', psi=psi, theta=theta, shape=20)
+            gen_data = pm.sample_prior_predictive(samples=5000)
+            assert gen_data['theta'].shape == (5000,)
+            assert gen_data['psi'].shape == (5000,)
+            assert gen_data['suppliers'].shape == (5000, 20)


Code review comments:
diff_hunk: @@ -964,6 +965,11 @@ def random(self, point=None, size=None):
         """
         mu, lam, alpha = draw_values([self.mu, self.lam, self.alpha],
                                      point=point, size=size)
+        print(mu.shape, lam.shape, alpha.shape)
path: pymc3/distributions/continuous.py
text: @user1:
dont forget to remove the `print`