Honor discard_tuned_samples during KeyboardInterrupt

pymc-devs · Jan 23, 2020 · 3a8503f · 3a8503f
1 parent b90de19
commit 3a8503f
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/pymc3/sampling.py b/pymc3/sampling.py
@@ -470,7 +470,7 @@ def sample(
         _log.info("Multiprocess sampling ({} chains in {} jobs)".format(chains, cores))
         _print_step_hierarchy(step)
         try:
-            trace = _mp_sample(**sample_args)
+            trace = _mp_sample(**sample_args, discard_tuned_samples=discard_tuned_samples)
         except pickle.PickleError:
             _log.warning("Could not pickle model, sampling singlethreaded.")
             _log.debug("Pickling error:", exec_info=True)
@@ -1243,6 +1243,7 @@ def _mp_sample(
     progressbar=True,
     trace=None,
     model=None,
+    discard_tuned_samples=True,
     **kwargs
 ):
     """Main iteration for multiprocess sampling.
@@ -1325,7 +1326,10 @@ def _mp_sample(
             raise
         return MultiTrace(traces)
     except KeyboardInterrupt:
-        traces, length = _choose_chains(traces, tune)
+        if discard_tuned_samples:
+            traces, length = _choose_chains(traces, tune)
+        else:
+            traces, length = _choose_chains(traces, 0)
         return MultiTrace(traces)[:length]
     finally:
         for trace in traces: