Require all step methods to return stats

The reason for this change is the resulting simplification of code,
including simpler branching and less type ambiguity.

Closes pymc-devs#6270

pymc/blocking.py

@@ -20,15 +20,18 @@
 from __future__ import annotations
 
 from functools import partial
-from typing import Callable, Dict, Generic, NamedTuple, TypeVar
+from typing import Any, Callable, Dict, Generic, List, NamedTuple, TypeVar
 
 import numpy as np
 
+from typing_extensions import TypeAlias
+
 __all__ = ["DictToArrayBijection"]
 
 
 T = TypeVar("T")
-PointType = Dict[str, np.ndarray]
+PointType: TypeAlias = Dict[str, np.ndarray]
+StatsType: TypeAlias = List[Dict[str, Any]]
 
 # `point_map_info` is a tuple of tuples containing `(name, shape, dtype)` for
 # each of the raveled variables.

pymc/sampling/mcmc.py

@@ -912,14 +912,10 @@ def _iter_sample(
                 step.iter_count = 0
             if i == tune:
                 step.stop_tuning()
-            if step.generates_stats:
-                point, stats = step.step(point)
-                strace.record(point, stats)
-                log_warning_stats(stats)
-                diverging = i > tune and stats and stats[0].get("diverging")
-            else:
-                point = step.step(point)
-                strace.record(point, [])
+            point, stats = step.step(point)
+            strace.record(point, stats)
+            log_warning_stats(stats)
+            diverging = i > tune and stats and stats[0].get("diverging")
             if callback is not None:
                 callback(
                     trace=strace,

pymc/sampling/parallel.py

@@ -173,7 +173,7 @@ def _start_loop(self):
 
             if draw < self._draws + self._tune:
                 try:
-                    point, stats = self._compute_point()
+                    point, stats = self._step_method.step(self._point)
                 except SamplingError as e:
                     e = ExceptionWithTraceback(e, e.__traceback__)
                     self._msg_pipe.send(("error", e))
@@ -191,14 +191,6 @@ def _start_loop(self):
             else:
                 raise ValueError("Unknown message " + msg[0])
 
-    def _compute_point(self):
-        if self._step_method.generates_stats:
-            point, stats = self._step_method.step(self._point)
-        else:
-            point = self._step_method.step(self._point)
-            stats = None
-        return point, stats
-
 
 def _run_process(*args):
     _Process(*args).run()

pymc/sampling/population.py

@@ -17,7 +17,7 @@
 import logging
 
 from copy import copy
-from typing import Iterator, List, Sequence, Union
+from typing import Iterator, List, Sequence, Tuple, Union
 
 import cloudpickle
 import numpy as np
@@ -31,7 +31,11 @@
 from pymc.model import modelcontext
 from pymc.stats.convergence import log_warning_stats
 from pymc.step_methods import CompoundStep
-from pymc.step_methods.arraystep import BlockedStep, PopulationArrayStepShared
+from pymc.step_methods.arraystep import (
+    BlockedStep,
+    PopulationArrayStepShared,
+    StatsType,
+)
 from pymc.util import RandomSeed
 
 __all__ = ()
@@ -224,7 +228,7 @@ def _run_secondary(c, stepper_dumps, secondary_end):
             _log.exception(f"ChainWalker{c}")
         return
 
-    def step(self, tune_stop: bool, population):
+    def step(self, tune_stop: bool, population) -> List[Tuple[PointType, StatsType]]:
         """Step the entire population of chains.
 
         Parameters
@@ -239,18 +243,18 @@ def step(self, tune_stop: bool, population):
         update : list
             List of (Point, stats) tuples for all chains
         """
-        updates = [None] * self.nchains
+        updates: List[Tuple[PointType, StatsType]] = []
         if self.is_parallelized:
             for c in range(self.nchains):
                 self._primary_ends[c].send((tune_stop, population))
             # Blockingly get the step outcomes
             for c in range(self.nchains):
-                updates[c] = self._primary_ends[c].recv()
+                updates.append(self._primary_ends[c].recv())
         else:
             for c in range(self.nchains):
                 if tune_stop:
                     self._steppers[c].stop_tuning()
-                updates[c] = self._steppers[c].step(population[c])
+                updates.append(self._steppers[c].step(population[c]))
         return updates
 
 
@@ -378,13 +382,9 @@ def _iter_population(
 
                 # apply the update to the points and record to the traces
                 for c, strace in enumerate(traces):
-                    if steppers[c].generates_stats:
-                        points[c], stats = updates[c]
-                        strace.record(points[c], stats)
-                        log_warning_stats(stats)
-                    else:
-                        points[c] = updates[c]
-                        strace.record(points[c])
+                    points[c], stats = updates[c]
+                    strace.record(points[c], stats)
+                    log_warning_stats(stats)
                 # yield the state of all chains in parallel
                 yield traces
     except KeyboardInterrupt:

pymc/step_methods/arraystep.py

@@ -14,22 +14,20 @@
 
 from abc import ABC, abstractmethod
 from enum import IntEnum, unique
-from typing import Dict, List, Tuple, TypeVar, Union
+from typing import Callable, Dict, List, Tuple, Union, cast
 
 import numpy as np
 
 from aesara.graph.basic import Variable
 from numpy.random import uniform
 
-from pymc.blocking import DictToArrayBijection, PointType, RaveledVars
+from pymc.blocking import DictToArrayBijection, PointType, RaveledVars, StatsType
 from pymc.model import modelcontext
 from pymc.step_methods.compound import CompoundStep
 from pymc.util import get_var_name
 
 __all__ = ["ArrayStep", "ArrayStepShared", "metrop_select", "Competence"]
 
-StatsType = TypeVar("StatsType")
-
 
 @unique
 class Competence(IntEnum):
@@ -49,7 +47,6 @@ class Competence(IntEnum):
 
 class BlockedStep(ABC):
 
-    generates_stats = False
     stats_dtypes: List[Dict[str, type]] = []
     vars: List[Variable] = []
 
@@ -103,7 +100,7 @@ def __getnewargs_ex__(self):
         return self.__newargs
 
     @abstractmethod
-    def step(point: PointType, *args, **kwargs) -> Union[PointType, Tuple[PointType, StatsType]]:
+    def step(self, point: PointType) -> Tuple[PointType, StatsType]:
         """Perform a single step of the sampler."""
 
     @staticmethod
@@ -146,35 +143,28 @@ def __init__(self, vars, fs, allvars=False, blocked=True):
         self.allvars = allvars
         self.blocked = blocked
 
-    def step(self, point: PointType):
+    def step(self, point: PointType) -> Tuple[PointType, StatsType]:
 
-        partial_funcs_and_point = [DictToArrayBijection.mapf(x, start_point=point) for x in self.fs]
+        partial_funcs_and_point: List[Union[Callable, PointType]] = [
+            DictToArrayBijection.mapf(x, start_point=point) for x in self.fs
+        ]
         if self.allvars:
             partial_funcs_and_point.append(point)
 
-        apoint = DictToArrayBijection.map({v.name: point[v.name] for v in self.vars})
-        step_res = self.astep(apoint, *partial_funcs_and_point)
-
-        if self.generates_stats:
-            apoint_new, stats = step_res
-        else:
-            apoint_new = step_res
+        var_dict = {cast(str, v.name): point[cast(str, v.name)] for v in self.vars}
+        apoint = DictToArrayBijection.map(var_dict)
+        apoint_new, stats = self.astep(apoint, *partial_funcs_and_point)
 
         if not isinstance(apoint_new, RaveledVars):
             # We assume that the mapping has stayed the same
             apoint_new = RaveledVars(apoint_new, apoint.point_map_info)
 
         point_new = DictToArrayBijection.rmap(apoint_new, start_point=point)
 
-        if self.generates_stats:
-            return point_new, stats
-
-        return point_new
+        return point_new, stats
 
     @abstractmethod
-    def astep(
-        self, apoint: RaveledVars, point: PointType, *args
-    ) -> Union[RaveledVars, Tuple[RaveledVars, StatsType]]:
+    def astep(self, apoint: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
         """Perform a single sample step in a raveled and concatenated parameter space."""
 
 
@@ -198,30 +188,27 @@ def __init__(self, vars, shared, blocked=True):
         self.shared = {get_var_name(var): shared for var, shared in shared.items()}
         self.blocked = blocked
 
-    def step(self, point):
+    def step(self, point: PointType) -> Tuple[PointType, StatsType]:
 
         for name, shared_var in self.shared.items():
             shared_var.set_value(point[name])
 
-        q = DictToArrayBijection.map({v.name: point[v.name] for v in self.vars})
+        var_dict = {cast(str, v.name): point[cast(str, v.name)] for v in self.vars}
+        q = DictToArrayBijection.map(var_dict)
 
-        step_res = self.astep(q)
-
-        if self.generates_stats:
-            apoint, stats = step_res
-        else:
-            apoint = step_res
+        apoint, stats = self.astep(q)
 
         if not isinstance(apoint, RaveledVars):
             # We assume that the mapping has stayed the same
             apoint = RaveledVars(apoint, q.point_map_info)
 
         new_point = DictToArrayBijection.rmap(apoint, start_point=point)
 
-        if self.generates_stats:
-            return new_point, stats
+        return new_point, stats
 
-        return new_point
+    @abstractmethod
+    def astep(self, q0: RaveledVars) -> Tuple[RaveledVars, StatsType]:
+        """Perform a single sample step in a raveled and concatenated parameter space."""
 
 
 class PopulationArrayStepShared(ArrayStepShared):
@@ -281,7 +268,7 @@ def __init__(
 
         super().__init__(vars, func._extra_vars_shared, blocked)
 
-    def step(self, point):
+    def step(self, point) -> Tuple[PointType, StatsType]:
         self._logp_dlogp_func._extra_are_set = True
         return super().step(point)
 

pymc/step_methods/compound.py

@@ -17,9 +17,12 @@
 
 @author: johnsalvatier
 """
-from collections import namedtuple
 
-import numpy as np
+
+from typing import Tuple
+
+from pymc.blocking import PointType
+from pymc.step_methods.arraystep import StatsType
 
 
 class CompoundStep:
@@ -28,36 +31,23 @@ class CompoundStep:
 
     def __init__(self, methods):
         self.methods = list(methods)
-        self.generates_stats = any(method.generates_stats for method in self.methods)
         self.stats_dtypes = []
         for method in self.methods:
-            if method.generates_stats:
-                self.stats_dtypes.extend(method.stats_dtypes)
+            self.stats_dtypes.extend(method.stats_dtypes)
         self.name = (
             f"Compound[{', '.join(getattr(m, 'name', 'UNNAMED_STEP') for m in self.methods)}]"
         )
 
-    def step(self, point):
-        if self.generates_stats:
-            states = []
-            for method in self.methods:
-                if method.generates_stats:
-                    point, state = method.step(point)
-                    states.extend(state)
-                else:
-                    point = method.step(point)
-            # Model logp can only be the logp of the _last_ state, if there is
-            # one. Pop all others (if dict), or set to np.nan (if namedtuple).
-            for state in states[:-1]:
-                if isinstance(state, dict):
-                    state.pop("model_logp", None)
-                elif isinstance(state, namedtuple):
-                    state = state._replace(logp=np.nan)
-            return point, states
-        else:
-            for method in self.methods:
-                point = method.step(point)
-            return point
+    def step(self, point) -> Tuple[PointType, StatsType]:
+        stats = []
+        for method in self.methods:
+            point, sts = method.step(point)
+            stats.extend(sts)
+        # Model logp can only be the logp of the _last_ stats,
+        # if there is one. Pop all others.
+        for sts in stats[:-1]:
+            sts.pop("model_logp", None)
+        return point, stats
 
     def stop_tuning(self):
         for method in self.methods:

pymc/step_methods/hmc/base_hmc.py

@@ -23,7 +23,7 @@
 import numpy as np
 
 from pymc.aesaraf import floatX
-from pymc.blocking import DictToArrayBijection, RaveledVars
+from pymc.blocking import DictToArrayBijection, RaveledVars, StatsType
 from pymc.exceptions import SamplingError
 from pymc.model import Point, modelcontext
 from pymc.stats.convergence import SamplerWarning, WarningType
@@ -157,7 +157,7 @@ def _hamiltonian_step(self, start, p0, step_size) -> HMCStepData:
         Subclasses must overwrite this abstract method and return an `HMCStepData` object.
         """
 
-    def astep(self, q0: RaveledVars) -> tuple[RaveledVars, list[dict[str, Any]]]:
+    def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
         """Perform a single HMC iteration."""
         perf_start = time.perf_counter()
         process_start = time.process_time()

pymc/step_methods/hmc/hmc.py

@@ -39,7 +39,6 @@ class HamiltonianMC(BaseHMC):
 
     name = "hmc"
     default_blocked = True
-    generates_stats = True
     stats_dtypes = [
         {
             "step_size": np.float64,

pymc/step_methods/hmc/nuts.py

@@ -97,7 +97,6 @@ class NUTS(BaseHMC):
     name = "nuts"
 
     default_blocked = True
-    generates_stats = True
     stats_dtypes = [
         {
             "depth": np.int64,