Fix parallel sampling (#1481)

* Fix parallel sampling

* Everything is good again

Author: ColCarroll

pymc3/sampling.py

@@ -2,7 +2,7 @@
 
 from joblib import Parallel, delayed
 from numpy.random import randint, seed
-from numpy import shape, asarray
+import numpy as np
 
 import pymc3 as pm
 from .backends.base import merge_traces, BaseTrace, MultiTrace
@@ -276,18 +276,26 @@ def _choose_backend(trace, chain, shortcuts=None, **kwds):
 
 
 def _make_parallel(arg, njobs):
-    if not shape(arg):
+    if not np.shape(arg):
         return [arg] * njobs
     return arg
 
 
+def _parallel_random_seed(random_seed, njobs):
+    if random_seed == -1 and njobs > 1:
+        max_int = np.iinfo(np.int32).max
+        return [randint(max_int) for _ in range(njobs)]
+    else:
+        return _make_parallel(random_seed, njobs)
+
+
 def _mp_sample(**kwargs):
     njobs = kwargs.pop('njobs')
     chain = kwargs.pop('chain')
     random_seed = kwargs.pop('random_seed')
     start = kwargs.pop('start')
 
-    rseed = _make_parallel(random_seed, njobs)
+    rseed = _parallel_random_seed(random_seed, njobs)
     start_vals = _make_parallel(start, njobs)
 
     chains = list(range(chain, chain + njobs))
@@ -365,4 +373,4 @@ def sample_ppc(trace, samples=None, model=None, vars=None, size=None, random_see
             ppc[var.name].append(var.distribution.random(point=param,
                                                          size=size))
 
-    return {k: asarray(v) for k, v in ppc.items()}
+    return {k: np.asarray(v) for k, v in ppc.items()}

pymc3/tests/test_sampling.py

@@ -1,3 +1,4 @@
+from itertools import combinations
 import numpy as np
 try:
     import unittest.mock as mock  # py3
@@ -33,6 +34,26 @@ def test_sample_does_not_set_seed(self):
                 random_numbers.append(np.random.random())
         self.assertEqual(random_numbers[0], random_numbers[1])
 
+    def test_parallel_sample_does_not_reuse_seed(self):
+        njobs = 4
+        random_numbers = []
+        draws = []
+        for _ in range(2):
+            np.random.seed(1)  # seeds in other processes don't effect main process
+            with self.model:
+                trace = pm.sample(100, njobs=njobs)
+            # numpy thread mentioned race condition.  might as well check none are equal
+            for first, second in combinations(range(njobs), 2):
+                first_chain = trace.get_values('x', chains=first)
+                second_chain = trace.get_values('x', chains=second)
+                self.assertFalse((first_chain == second_chain).all())
+            draws.append(trace.get_values('x'))
+            random_numbers.append(np.random.random())
+
+        # Make sure future random processes aren't effected by this
+        self.assertEqual(*random_numbers)
+        self.assertTrue((draws[0] == draws[1]).all())
+
     def test_sample(self):
         test_njobs = [1]
         with self.model: