Truncate retrieved arrays to chain length

mcbackend/__init__.py

@@ -20,4 +20,4 @@
     pass
 
 
-__version__ = "0.1.1"
+__version__ = "0.1.2"

mcbackend/core.py

@@ -188,9 +188,13 @@ def to_inferencedata(self, **kwargs) -> InferenceData:
         posterior = collections.defaultdict(list)
         sample_stats = collections.defaultdict(list)
         for c, chain in enumerate(chains):
+            # Every retrieved array is shortened to the previously determined chain length.
+            # This is needed for database backends which may get inserts inbetween.
+            clen = chain_lengths[chain.cid]
+
             # Obtain a mask by which draws can be split into warmup/posterior
             if "tune" in chain.sample_stats:
-                tune = chain.get_stats("tune").astype(bool)
+                tune = chain.get_stats("tune")[:clen].astype(bool)
             else:
                 if c == 0:
                     _log.warning(
@@ -200,12 +204,12 @@ def to_inferencedata(self, **kwargs) -> InferenceData:
 
             # Split all variables draws into warmup/posterior
             for var in variables:
-                draws = chain.get_draws(var.name)
+                draws = chain.get_draws(var.name)[:clen]
                 warmup_posterior[var.name].append(draws[tune])
                 posterior[var.name].append(draws[~tune])
             # Same for sample stats
             for svar in self.meta.sample_stats:
-                stats = chain.get_stats(svar.name)
+                stats = chain.get_stats(svar.name)[:clen]
                 warmup_sample_stats[svar.name].append(stats[tune])
                 sample_stats[svar.name].append(stats[~tune])
 

requirements-dev.txt

@@ -1,8 +1,10 @@
 arviz
 clickhouse-driver
 flake8
-pymc>=4.0.1
+pymc==4.1.4
 pytest
 pytest-cov
+# Temporary pin to fix Aesara import
+setuptools>=48.0.0
 twine
 wheel