use chain order rather than logical index with _DefaultTrace.insert

pymc-devs · Mar 9, 2021 · 023fdcd · 023fdcd
1 parent 4d535a9
commit 023fdcd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pymc3/backends/arviz.py b/pymc3/backends/arviz.py
@@ -206,12 +206,12 @@ def _extract_log_likelihood(self, trace):
             ]
         log_likelihood_dict = _DefaultTrace(len(trace.chains))
         for var, log_like_fun in cached:
-            for chain in trace.chains:
+            for k, chain in enumerate(trace.chains):
                 log_like_chain = [
                     self.log_likelihood_vals_point(point, var, log_like_fun)
                     for point in trace.points([chain])
                 ]
-                log_likelihood_dict.insert(var.name, np.stack(log_like_chain), chain)
+                log_likelihood_dict.insert(var.name, np.stack(log_like_chain), k)
         return log_likelihood_dict.trace_dict
 
     @requires("trace")