Made sample_shape same across all contexts in draw_values

RELEASE-NOTES.md

@@ -1,5 +1,10 @@
 # Release Notes
 
+## PyMC3 3.10.x
+
+### Maintenance
+- Make `sample_shape` same across all contexts in `draw_values` (see [#4305](https://github.com/pymc-devs/pymc3/pull/4305)).
+
 ## PyMC3 3.10.0 (7 December 2020)
 
 This is a major release with many exciting new features. The biggest change is that we now rely on our own fork of [Theano-PyMC](https://github.com/pymc-devs/Theano-PyMC). This is in line with our [big announcement about our commitment to PyMC3 and Theano](https://pymc-devs.medium.com/the-future-of-pymc3-or-theano-is-dead-long-live-theano-d8005f8a0e9b).

pymc3/distributions/distribution.py

@@ -690,6 +690,7 @@ def draw_values(params, point=None, size=None):
     """
     # The following check intercepts and redirects calls to
     # draw_values in the context of sample_posterior_predictive
+    size = to_tuple(size)
     ppc_sampler = vectorized_ppc.get(None)
     if ppc_sampler is not None:
         # this is being done inside new, vectorized sample_posterior_predictive

pymc3/distributions/simulator.py

@@ -18,7 +18,7 @@
 
 from scipy.spatial import cKDTree
 
-from pymc3.distributions.distribution import NoDistribution, draw_values
+from pymc3.distributions.distribution import NoDistribution, draw_values, to_tuple
 
 __all__ = ["Simulator"]
 
@@ -114,11 +114,12 @@ def random(self, point=None, size=None):
         -------
         array
         """
+        size = to_tuple(size)
         params = draw_values([*self.params], point=point, size=size)
-        if size is None:
+        if len(size) == 0:
             return self.function(*params)
         else:
-            return np.array([self.function(*params) for _ in range(size)])
+            return np.array([self.function(*params) for _ in range(size[0])])
 
     def _str_repr(self, name=None, dist=None, formatting="plain"):
         if dist is None:

pymc3/tests/test_distributions_random.py

@@ -1664,6 +1664,10 @@ def test_issue_3758(self):
         for var in "abcd":
             assert not np.isnan(np.std(samples[var]))
 
+        for var in "bcd":
+            std = np.std(samples[var] - samples["a"])
+            np.testing.assert_allclose(std, 1, rtol=1e-2)
+
     def test_issue_3829(self):
         with pm.Model() as model:
             x = pm.MvNormal("x", mu=np.zeros(5), cov=np.eye(5), shape=(2, 5))