Emit "tune" from all stepmethods

Author: michaelosthege

pymc/step_methods/compound.py

@@ -262,6 +262,16 @@ def flatten_steps(step: Union[BlockedStep, CompoundStep]) -> List[BlockedStep]:
     return steps
 
 
+def check_step_emits_tune(step: Union[CompoundStep, BlockedStep]):
+    if isinstance(step, BlockedStep) and "tune" not in step.stats_dtypes_shapes:
+        raise TypeError(f"{type(step)} does not emit the required 'tune' stat.")
+    elif isinstance(step, CompoundStep):
+        for sstep in step.methods:
+            if "tune" not in sstep.stats_dtypes_shapes:
+                raise TypeError(f"{type(sstep)} does not emit the required 'tune' stat.")
+    return
+
+
 class StatsBijection:
     """Map between a `list` of stats to `dict` of stats."""
 

pymc/step_methods/metropolis.py

@@ -458,9 +458,16 @@ class BinaryGibbsMetropolis(ArrayStep):
 
     name = "binary_gibbs_metropolis"
 
+    stats_dtypes_shapes = {
+        "tune": (bool, []),
+    }
+
     def __init__(self, vars, order="random", transit_p=0.8, model=None):
         model = pm.modelcontext(model)
 
+        # Doesn't actually tune, but it's required to emit a sampler stat
+        # that indicates whether a draw was done in a tuning phase.
+        self.tune = True
         # transition probabilities
         self.transit_p = transit_p
 
@@ -483,6 +490,11 @@ def __init__(self, vars, order="random", transit_p=0.8, model=None):
 
         super().__init__(vars, [model.compile_logp()])
 
+    def reset_tuning(self):
+        # There are no tuning parameters in this step method.
+        self.tune = False
+        return
+
     def astep(self, apoint: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
         logp: Callable[[RaveledVars], np.ndarray] = args[0]
         order = self.order
@@ -503,7 +515,10 @@ def astep(self, apoint: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
                 if accepted:
                     logp_curr = logp_prop
 
-        return q, []
+        stats = {
+            "tune": self.tune,
+        }
+        return q, [stats]
 
     @staticmethod
     def competence(var):
@@ -543,6 +558,10 @@ class CategoricalGibbsMetropolis(ArrayStep):
 
     name = "categorical_gibbs_metropolis"
 
+    stats_dtypes_shapes = {
+        "tune": (bool, []),
+    }
+
     def __init__(self, vars, proposal="uniform", order="random", model=None):
         model = pm.modelcontext(model)
 
@@ -593,8 +612,17 @@ def __init__(self, vars, proposal="uniform", order="random", model=None):
         else:
             raise ValueError("Argument 'proposal' should either be 'uniform' or 'proportional'")
 
+        # Doesn't actually tune, but it's required to emit a sampler stat
+        # that indicates whether a draw was done in a tuning phase.
+        self.tune = True
+
         super().__init__(vars, [model.compile_logp()])
 
+    def reset_tuning(self):
+        # There are no tuning parameters in this step method.
+        self.tune = False
+        return
+
     def astep_unif(self, apoint: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
         logp = args[0]
         point_map_info = apoint.point_map_info
@@ -614,7 +642,10 @@ def astep_unif(self, apoint: RaveledVars, *args) -> Tuple[RaveledVars, StatsType
             if accepted:
                 logp_curr = logp_prop
 
-        return q, []
+        stats = {
+            "tune": self.tune,
+        }
+        return q, [stats]
 
     def astep_prop(self, apoint: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
         logp = args[0]

tests/step_methods/test_compound.py

@@ -26,6 +26,7 @@
     Slice,
 )
 from pymc.step_methods.compound import (
+    BlockedStep,
     StatsBijection,
     flatten_steps,
     get_stats_dtypes_shapes_from_steps,
@@ -36,6 +37,16 @@
 from tests.models import simple_2model_continuous
 
 
+def test_all_stepmethods_emit_tune_stat():
+    attrs = [getattr(pm.step_methods, n) for n in dir(pm.step_methods)]
+    step_types = [
+        attr for attr in attrs if isinstance(attr, type) and issubclass(attr, BlockedStep)
+    ]
+    assert len(step_types) > 5
+    for cls in step_types:
+        assert "tune" in cls.stats_dtypes_shapes
+
+
 class TestCompoundStep:
     samplers = (Metropolis, Slice, HamiltonianMC, NUTS, DEMetropolis)