Assert warmup/posterior lengths in StepMethodTesteer

Author: michaelosthege

tests/helpers.py

@@ -133,7 +133,7 @@ def check_stat_dtype(self, step, idata):
                     continue
                 assert idata.sample_stats[stat].dtype == np.dtype(dtype)
 
-    def step_continuous(self, step_fn, draws):
+    def step_continuous(self, step_fn, draws, chains=1, tune=1000):
         start, model, (mu, C) = mv_simple()
         unc = np.diag(C) ** 0.5
         check = (("x", np.mean, mu, unc / 10), ("x", np.std, unc, unc / 10))
@@ -143,14 +143,19 @@ def step_continuous(self, step_fn, draws):
             with warnings.catch_warnings():
                 warnings.filterwarnings("ignore", "More chains .* than draws .*", UserWarning)
                 idata = pm.sample(
-                    tune=1000,
+                    tune=tune,
                     draws=draws,
-                    chains=1,
+                    chains=chains,
                     step=step,
                     initvals=start,
                     model=model,
                     random_seed=1,
+                    discard_tuned_samples=False,
                 )
+            assert idata.warmup_posterior.sizes["chain"] == chains
+            assert idata.warmup_posterior.sizes["draw"] == tune
+            assert idata.posterior.sizes["chain"] == chains
+            assert idata.posterior.sizes["draw"] == draws
             self.check_stat(check, idata, step.__class__.__name__)
             self.check_stat_dtype(idata, step)