Check chain lengths when merging.

Raise ValueError for mismatched chains.
Added test to verify.

Author: rpgoldman

Diff for: pymc3/backends/base.py

@@ -5,6 +5,7 @@
 """
 import itertools as itl
 import logging
+from typing import List
 
 import numpy as np
 import warnings
@@ -92,7 +93,8 @@ def _set_sampler_vars(self, sampler_vars):
 
         self.sampler_vars = sampler_vars
 
-    def setup(self, draws, chain, sampler_vars=None):
+    # pylint: disable=unused-argument
+    def setup(self, draws, chain, sampler_vars=None) -> None: 
         """Perform chain-specific setup.
 
         Parameters
@@ -542,7 +544,7 @@ def points(self, chains=None):
         return itl.chain.from_iterable(self._straces[chain] for chain in chains)
 
 
-def merge_traces(mtraces):
+def merge_traces(mtraces: List[MultiTrace]) -> MultiTrace:
     """Merge MultiTrace objects.
 
     Parameters
@@ -552,17 +554,26 @@ def merge_traces(mtraces):
 
     Raises
     ------
-    A ValueError is raised if any traces have overlapping chain numbers.
+    A ValueError is raised if any traces have overlapping chain numbers,
+    or if chains are of different lengths.
 
     Returns
     -------
     A MultiTrace instance with merged chains
     """
+    if len(mtraces) == 0:
+        raise ValueError("Cannot merge an empty set of traces.")
     base_mtrace = mtraces[0]
+    chain_len = len(base_mtrace)
+    # check base trace
+    if any((len(st) != chain_len for _, st in base_mtrace._straces.items())):
+        raise ValueError("Chains are of different lengths.")
     for new_mtrace in mtraces[1:]:
         for new_chain, strace in new_mtrace._straces.items():
             if new_chain in base_mtrace._straces:
                 raise ValueError("Chains are not unique.")
+            if len(strace) != chain_len:
+                raise ValueError("Chains are of different lengths.")
             base_mtrace._straces[new_chain] = strace
     base_mtrace._report = merge_reports([trace.report for trace in mtraces])
     return base_mtrace

Diff for: pymc3/tests/test_ndarray_backend.py

@@ -112,6 +112,30 @@ def test_multitrace_nonunique(self):
         with pytest.raises(ValueError):
             base.MultiTrace([self.strace0, self.strace1])
 
+    def test_merge_traces_no_traces(self):
+        with pytest.raises(ValueError):
+            base.merge_traces([])
+
+    def test_merge_traces_diff_lengths(self):
+        with self.model:
+            strace0 = self.backend(self.name)
+            strace0.setup(self.draws, 1)
+            for i in range(self.draws):
+                strace0.record(self.test_point)
+            strace0.close()
+        mtrace0 = base.MultiTrace([self.strace0])
+
+        with self.model:
+            strace1 = self.backend(self.name)
+            strace1.setup(2 * self.draws, 1)
+            for i in range(2 * self.draws):
+                strace1.record(self.test_point)
+            strace1.close()
+        mtrace1 = base.MultiTrace([strace1])
+
+        with pytest.raises(ValueError):
+            base.merge_traces([mtrace0, mtrace1])
+
     def test_merge_traces_nonunique(self):
         mtrace0 = base.MultiTrace([self.strace0])
         mtrace1 = base.MultiTrace([self.strace1])