Adapt to updates in upstream behavior

Author: lucianopaz

pymc/backends/ndarray.py

@@ -232,8 +232,8 @@ def completed_draws_and_divergences(self, chain_specific: bool = True) -> tuple[
         for sampler_stats in self._stats:
             for key, data in sampler_stats.items():
                 if "divergence" in key:
-                    divergent_draws += np.asarray(data)
-        divergences = sum(divergent_draws > 0)
+                    divergent_draws += np.asarray(data[: len(self)])
+        divergences = int(sum(divergent_draws > 0, start=0))
         return len(self), divergences
 
 

pymc/sampling/mcmc.py

@@ -1155,7 +1155,9 @@ def _sample_many(
                 *trace.completed_draws_and_divergences(chain_specific=True)
             )
             progress_manager._progress.update(
-                progress_manager.tasks[i],
+                progress_manager.tasks[0]
+                if progress_manager.combined_progress
+                else progress_manager.tasks[i],
                 draws=progress_manager.completed_draws
                 if progress_manager.combined_progress
                 else progress_manager.draws,

pymc/step_methods/hmc/base_hmc.py

@@ -69,6 +69,7 @@ class BaseHMCState(StepMethodState):
     tune: bool
     potential: PotentialState
     _num_divs_sample: int
+    divergences: int
 
 
 class BaseHMC(GradientSharedStep):