Create traces with init_traces helper function

Author: michaelosthege

pymc/backends/__init__.py

@@ -61,12 +61,15 @@
 
 """
 from copy import copy
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Sequence, Union
+
+import numpy as np
 
 from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
-from pymc.backends.base import BaseTrace
+from pymc.backends.base import BaseTrace, IBaseTrace
 from pymc.backends.ndarray import NDArray
 from pymc.model import Model
+from pymc.step_methods.compound import BlockedStep, CompoundStep
 
 __all__ = ["to_inference_data", "predictions_to_inference_data"]
 
@@ -92,3 +95,26 @@ def _init_trace(
 
     strace.setup(expected_length, chain_number, stats_dtypes)
     return strace
+
+
+def init_traces(
+    *,
+    backend: Optional[BaseTrace],
+    chains: int,
+    expected_length: int,
+    step: Union[BlockedStep, CompoundStep],
+    var_dtypes: Dict[str, np.dtype],
+    var_shapes: Dict[str, Sequence[int]],
+    model: Model,
+) -> Sequence[IBaseTrace]:
+    """Initializes a trace recorder for each chain."""
+    return [
+        _init_trace(
+            expected_length=expected_length,
+            stats_dtypes=step.stats_dtypes,
+            chain_number=chain_number,
+            trace=backend,
+            model=model,
+        )
+        for chain_number in range(chains)
+    ]

pymc/sampling/mcmc.py

@@ -32,7 +32,7 @@
 
 import pymc as pm
 
-from pymc.backends import _init_trace
+from pymc.backends import init_traces
 from pymc.backends.base import BaseTrace, IBaseTrace, MultiTrace, _choose_chains
 from pymc.blocking import DictToArrayBijection
 from pymc.exceptions import SamplingError
@@ -486,21 +486,21 @@ def sample(
         initial_points = [ipfn(seed) for ipfn, seed in zip(ipfns, random_seed_list)]
 
     # One final check that shapes and logps at the starting points are okay.
+    ip: Dict[str, np.ndarray]
     for ip in initial_points:
         model.check_start_vals(ip)
         _check_start_shape(model, ip)
 
     # Create trace backends for each chain
-    traces = [
-        _init_trace(
-            expected_length=draws + tune,
-            stats_dtypes=step.stats_dtypes,
-            chain_number=chain_number,
-            trace=trace,
-            model=model,
-        )
-        for chain_number in range(chains)
-    ]
+    traces = init_traces(
+        backend=trace,
+        chains=chains,
+        expected_length=draws + tune,
+        step=step,
+        var_dtypes={vn: v.dtype for vn, v in ip.items()},
+        var_shapes={vn: v.shape for vn, v in ip.items()},
+        model=model,
+    )
 
     sample_args = {
         "draws": draws,