Don't store total divergences in NUTS stats

Author: jessegrabowski

pymc/step_methods/hmc/base_hmc.py

@@ -275,7 +275,6 @@ 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/nuts.py

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