assume supports_sample_stats=True

Author: williambdean

pymc/backends/base.py

@@ -52,8 +52,6 @@ class BaseTrace(ABC):
         use different test point that might be with changed variables shapes
     """
 
-    supports_sampler_stats = False
-
     def __init__(self, name, model=None, vars=None, test_point=None):
         self.name = name
 
@@ -88,9 +86,6 @@ def _add_warnings(self, warnings):
     # Sampling methods
 
     def _set_sampler_vars(self, sampler_vars):
-        if sampler_vars is not None and not self.supports_sampler_stats:
-            raise ValueError("Backend does not support sampler stats.")
-
         if self._is_base_setup and self.sampler_vars != sampler_vars:
             raise ValueError("Can't change sampler_vars")
 
@@ -117,9 +112,7 @@ def setup(self, draws, chain, sampler_vars=None) -> None:
         chain: int
             Chain number
         sampler_vars: list of dictionaries (name -> dtype), optional
-            Diagnostics / statistics for each sampler. Before passing this
-            to a backend, you should check, that the `supports_sampler_state`
-            flag is set.
+            Diagnostics / statistics for each sampler
         """
         self._set_sampler_vars(sampler_vars)
         self._is_base_setup = True
@@ -190,9 +183,6 @@ def get_sampler_stats(self, stat_name, sampler_idx=None, burn=0, thin=1):
         a numpy array of shape (m, n), where `m` is the number of
         such samplers, and `n` is the number of samples.
         """
-        if not self.supports_sampler_stats:
-            raise ValueError("This backend does not support sampler stats")
-
         if sampler_idx is not None:
             return self._get_sampler_stats(stat_name, sampler_idx, burn, thin)
 
@@ -232,13 +222,11 @@ def point(self, idx):
 
     @property
     def stat_names(self):
-        if self.supports_sampler_stats:
-            names = set()
-            for vars in self.sampler_vars or []:
-                names.update(vars.keys())
-            return names
-        else:
-            return set()
+        names = set()
+        for vars in self.sampler_vars or []:
+            names.update(vars.keys())
+
+        return names
 
 
 class MultiTrace:
@@ -356,7 +344,7 @@ def __getitem__(self, idx):
             return self.get_sampler_stats(var, burn=burn, thin=thin)
         raise KeyError("Unknown variable %s" % var)
 
-    _attrs = {"_straces", "varnames", "chains", "stat_names", "supports_sampler_stats", "_report"}
+    _attrs = {"_straces", "varnames", "chains", "stat_names", "_report"}
 
     def __getattr__(self, name):
         # Avoid infinite recursion when called before __init__

pymc/backends/ndarray.py

@@ -41,8 +41,6 @@ class NDArray(base.BaseTrace):
         `model.unobserved_RVs` is used.
     """
 
-    supports_sampler_stats = True
-
     def __init__(self, name=None, model=None, vars=None, test_point=None):
         super().__init__(name, model, vars, test_point)
         self.draw_idx = 0

pymc/sampling.py

@@ -1048,11 +1048,8 @@ def _iter_sample(
                 step = stop_tuning(step)
             if step.generates_stats:
                 point, stats = step.step(point)
-                if strace.supports_sampler_stats:
-                    strace.record(point, stats)
-                    diverging = i > tune and stats and stats[0].get("diverging")
-                else:
-                    strace.record(point)
+                strace.record(point, stats)
+                diverging = i > tune and stats and stats[0].get("diverging")
             else:
                 point = step.step(point)
                 strace.record(point)
@@ -1359,10 +1356,7 @@ def _iter_population(
                 for c, strace in enumerate(traces):
                     if steppers[c].generates_stats:
                         points[c], stats = updates[c]
-                        if strace.supports_sampler_stats:
-                            strace.record(points[c], stats)
-                        else:
-                            strace.record(points[c])
+                        strace.record(points[c], stats)
                     else:
                         points[c] = updates[c]
                         strace.record(points[c])
@@ -1428,7 +1422,7 @@ def _init_trace(
     else:
         strace = _choose_backend(None, model=model)
 
-    if step.generates_stats and strace.supports_sampler_stats:
+    if step.generates_stats:
         strace.setup(expected_length, chain_number, step.stats_dtypes)
     else:
         strace.setup(expected_length, chain_number)
@@ -1520,7 +1514,7 @@ def _mp_sample(
             with sampler:
                 for draw in sampler:
                     strace = traces[draw.chain]
-                    if strace.supports_sampler_stats and draw.stats is not None:
+                    if draw.stats is not None:
                         strace.record(draw.point, draw.stats)
                     else:
                         strace.record(draw.point)

pymc/tests/backends/fixtures.py

@@ -51,7 +51,6 @@ def setup_method(self):
         if not hasattr(self, "sampler_vars"):
             self.sampler_vars = None
         if self.sampler_vars is not None:
-            assert self.strace.supports_sampler_stats
             self.strace.setup(self.draws, self.chain, self.sampler_vars)
         else:
             self.strace.setup(self.draws, self.chain)
@@ -110,11 +109,7 @@ def test_bad_dtype(self):
         with pytest.raises((ValueError, TypeError)):
             strace.setup(self.draws, self.chain, bad_vars)
         strace.setup(self.draws, self.chain, good_vars)
-        if strace.supports_sampler_stats:
-            assert strace.stat_names == {"a"}
-        else:
-            with pytest.raises((ValueError, TypeError)):
-                strace.setup(self.draws, self.chain, good_vars)
+        assert strace.stat_names == {"a"}
 
     def teardown_method(self):
         if self.name is not None: