proper handling of chain_idx in sample() (#4495)

* don't add chain offset in loop already taking care of this

* don't default to chain 0 when computing sampler stats; use provided chain_idx

* adding test script for chain_idx in sample()

* marking test as xfail for now

* take care of chain indices in sample_posterior_predictive

* update test to include sample_posterior_predictive

* use reproducable order of ppc samples wrt multiple chains

Author: Spaak

pymc3/sampling.py

@@ -600,7 +600,7 @@ def sample(
     # count the number of tune/draw iterations that happened
     # ideally via the "tune" statistic, but not all samplers record it!
     if "tune" in trace.stat_names:
-        stat = trace.get_sampler_stats("tune", chains=0)
+        stat = trace.get_sampler_stats("tune", chains=chain_idx)
         # when CompoundStep is used, the stat is 2 dimensional!
         if len(stat.shape) == 2:
             stat = stat[:, 0]
@@ -1453,9 +1453,9 @@ def _mp_sample(
         # dict does not contain all parameters
         update_start_vals(start[idx - chain], model.test_point, model)
         if step.generates_stats and strace.supports_sampler_stats:
-            strace.setup(draws + tune, idx + chain, step.stats_dtypes)
+            strace.setup(draws + tune, idx, step.stats_dtypes)
         else:
-            strace.setup(draws + tune, idx + chain)
+            strace.setup(draws + tune, idx)
         traces.append(strace)
 
     sampler = ps.ParallelSampler(
@@ -1716,12 +1716,19 @@ def sample_posterior_predictive(
 
     ppc_trace_t = _DefaultTrace(samples)
     try:
+        if hasattr(_trace, "_straces"):
+            # trace dict is unordered, but we want to return ppc samples in
+            # a predictable ordering, so sort the chain indices
+            chain_idx_mapping = sorted(_trace._straces.keys())
         for idx in indices:
             if nchain > 1:
                 # the trace object will either be a MultiTrace (and have _straces)...
                 if hasattr(_trace, "_straces"):
                     chain_idx, point_idx = np.divmod(idx, len_trace)
-                    param = cast(MultiTrace, _trace)._straces[chain_idx % nchain].point(point_idx)
+                    chain_idx = chain_idx % nchain
+                    # chain indices might not always start at 0, convert to proper index
+                    chain_idx = chain_idx_mapping[chain_idx]
+                    param = cast(MultiTrace, _trace)._straces[chain_idx].point(point_idx)
                 # ... or a PointList
                 else:
                     param = cast(PointList, _trace)[idx % (len_trace * nchain)]

pymc3/tests/test_sampling.py

@@ -321,6 +321,21 @@ def test_partial_trace_sample():
         trace = pm.sample(trace=[a])
 
 
+@pytest.mark.xfail
+def test_chain_idx():
+    # see https://github.com/pymc-devs/pymc3/issues/4469
+    with pm.Model():
+        mu = pm.Normal("mu")
+        x = pm.Normal("x", mu=mu, sigma=1, observed=np.asarray(3))
+        # note draws-tune must be >100 AND we need an observed RV for this to properly
+        # trigger convergence checks, which is one particular case in which this failed
+        # before
+        trace = pm.sample(draws=150, tune=10, chain_idx=1)
+
+        ppc = pm.sample_posterior_predictive(trace)
+        ppc = pm.sample_posterior_predictive(trace, keep_size=True)
+
+
 @pytest.mark.parametrize(
     "n_points, tune, expected_length, expected_n_traces",
     [