Track divergences in HMC stats

Author: jessegrabowski

pymc/step_methods/hmc/base_hmc.py

@@ -275,6 +275,7 @@ def astep(self, q0: RaveledVars) -> tuple[RaveledVars, StatsType]:
         stats: dict[str, Any] = {
             "tune": self.tune,
             "diverging": diverging,
+            "divergences": self.divergences,
             "perf_counter_diff": perf_end - perf_start,
             "process_time_diff": process_end - process_start,
             "perf_counter_start": perf_start,

pymc/step_methods/hmc/hmc.py

@@ -221,12 +221,12 @@ def _progressbar_config(n_chains=1):
 
         return columns, stats
 
-    def _make_progressbar_update_functions(self):
+    @staticmethod
+    def _make_progressbar_update_functions():
         def update_stats(stats):
-            divergences = self.divergences
             return {key: stats[key] for key in ("n_steps",)} | {
-                "failing": divergences > 0,
-                "divergences": divergences,
+                "failing": stats["divergences"] > 0,
+                "divergences": stats["divergences"],
             }
 
         return (update_stats,)

pymc/step_methods/hmc/nuts.py

@@ -115,6 +115,7 @@ class NUTS(BaseHMC):
         "step_size_bar": (np.float64, []),
         "tree_size": (np.float64, []),
         "diverging": (bool, []),
+        "divergences": (int, []),
         "energy_error": (np.float64, []),
         "energy": (np.float64, []),
         "max_energy_error": (np.float64, []),
@@ -247,12 +248,12 @@ def _progressbar_config(n_chains=1):
 
         return columns, stats
 
-    def _make_update_stats_functions(self):
+    @staticmethod
+    def _make_update_stats_functions():
         def update_stats(stats):
-            divergences = self.divergences
             return {key: stats[key] for key in ("step_size", "tree_size")} | {
-                "failing": divergences > 0,
-                "divergences": divergences,
+                "failing": stats["divergences"] > 0,
+                "divergences": stats["divergences"],
             }
 
         return (update_stats,)