Call run_convergence_checks after sampling with blackjax

jessegrabowski · jessegrabowski · commit e28c71ea6cdf · 2023-12-06T20:11:48.000+01:00
diff --git a/pymc/sampling/mcmc.py b/pymc/sampling/mcmc.py
@@ -380,6 +380,9 @@ def _sample_external_nuts(
             idata_kwargs=idata_kwargs,
             **nuts_sampler_kwargs,
         )
+        warns = run_convergence_checks(idata, model)
+        log_warnings(warns)
+
         return idata
 
     else:

Original file line number	Diff line number	Diff line change
`@@ -380,6 +380,9 @@ def _sample_external_nuts(`
`380`	`380`	`idata_kwargs=idata_kwargs,`
`381`	`381`	`**nuts_sampler_kwargs,`
`382`	`382`	`)`
	`383`	`+ warns = run_convergence_checks(idata, model)`
	`384`	`+ log_warnings(warns)`
	`385`	`+`
`383`	`386`	`return idata`
`384`	`387`
`385`	`388`	`else:`