Modify internal sampling functions from BaseTrace to IBaseTrace

Author: michaelosthege

pymc/sampling/mcmc.py

@@ -33,7 +33,7 @@
 import pymc as pm
 
 from pymc.backends import _init_trace
-from pymc.backends.base import BaseTrace, MultiTrace, _choose_chains
+from pymc.backends.base import BaseTrace, IBaseTrace, MultiTrace, _choose_chains
 from pymc.blocking import DictToArrayBijection
 from pymc.exceptions import SamplingError
 from pymc.initial_point import PointType, StartDict, make_initial_point_fns_per_chain
@@ -71,7 +71,7 @@
 class SamplingIteratorCallback(Protocol):
     """Signature of the callable that may be passed to `pm.sample(callable=...)`."""
 
-    def __call__(self, trace: BaseTrace, draw: Draw):
+    def __call__(self, trace: IBaseTrace, draw: Draw):
         pass
 
 
@@ -657,7 +657,7 @@ def _sample_many(
     *,
     draws: int,
     chains: int,
-    traces: Sequence[BaseTrace],
+    traces: Sequence[IBaseTrace],
     start: Sequence[PointType],
     random_seed: Optional[Sequence[RandomSeed]],
     step: Step,
@@ -701,7 +701,7 @@ def _sample(
     start: PointType,
     draws: int,
     step: Step,
-    trace: BaseTrace,
+    trace: IBaseTrace,
     tune: int,
     model: Optional[Model] = None,
     callback=None,
@@ -726,8 +726,8 @@ def _sample(
         The number of samples to draw
     step : function
         Step function
-    trace : backend, optional
-        A backend instance.
+    trace
+        A chain backend to record draws and stats.
     tune : int
         Number of iterations to tune.
     model : Model (optional if in ``with`` context)
@@ -767,7 +767,7 @@ def _iter_sample(
     draws: int,
     step: Step,
     start: PointType,
-    trace: BaseTrace,
+    trace: IBaseTrace,
     chain: int = 0,
     tune: int = 0,
     model: Optional[Model] = None,
@@ -785,8 +785,8 @@ def _iter_sample(
     start : dict
         Starting point in parameter space (or partial point).
         Must contain numeric (transformed) initial values for all (transformed) free variables.
-    trace : backend
-        A backend instance.
+    trace
+        A chain backend to record draws and stats.
     chain : int, optional
         Chain number used to store sample in backend.
     tune : int, optional
@@ -852,7 +852,7 @@ def _mp_sample(
     random_seed: Sequence[RandomSeed],
     start: Sequence[PointType],
     progressbar: bool = True,
-    traces: Sequence[BaseTrace],
+    traces: Sequence[IBaseTrace],
     model: Optional[Model] = None,
     callback: Optional[SamplingIteratorCallback] = None,
     mp_ctx=None,
@@ -879,9 +879,8 @@ def _mp_sample(
         Dicts must contain numeric (transformed) initial values for all (transformed) free variables.
     progressbar : bool
         Whether or not to display a progress bar in the command line.
-    trace : BaseTrace, optional
-        A backend instance, or None.
-        If None, the NDArray backend is used.
+    traces
+        Recording backends for each chain.
     model : Model (optional if in ``with`` context)
     callback
         A function which gets called for every sample from the trace of a chain. The function is