Fix return_inferencedata warnings in tests

Author: maresb

benchmarks/benchmarks/benchmarks.py

@@ -98,6 +98,7 @@ def time_overhead_sample(self, step):
                 random_seed=1,
                 progressbar=False,
                 compute_convergence_checks=False,
+                return_inferencedata=False,
             )
 
 
@@ -150,13 +151,23 @@ def time_drug_evaluation(self):
                 "effect size", diff_of_means / np.sqrt((group1_std ** 2 + group2_std ** 2) / 2)
             )
             pm.sample(
-                draws=20000, cores=4, chains=4, progressbar=False, compute_convergence_checks=False
+                draws=20000,
+                cores=4,
+                chains=4,
+                progressbar=False,
+                compute_convergence_checks=False,
+                return_inferencedata=False,
             )
 
     def time_glm_hierarchical(self):
         with glm_hierarchical_model():
             pm.sample(
-                draws=20000, cores=4, chains=4, progressbar=False, compute_convergence_checks=False
+                draws=20000,
+                cores=4,
+                chains=4,
+                progressbar=False,
+                compute_convergence_checks=False,
+                return_inferencedata=False,
             )
 
 
@@ -190,6 +201,7 @@ def track_glm_hierarchical_ess(self, init):
                 random_seed=100,
                 progressbar=False,
                 compute_convergence_checks=False,
+                return_inferencedata=False,
             )
             tot = time.time() - t0
         ess = float(az.ess(trace, var_names=["mu_a"])["mu_a"].values)
@@ -212,6 +224,7 @@ def track_marginal_mixture_model_ess(self, init):
                 random_seed=100,
                 progressbar=False,
                 compute_convergence_checks=False,
+                return_inferencedata=False,
             )
             tot = time.time() - t0
         ess = az.ess(trace, var_names=["mu"])["mu"].values.min()  # worst case
@@ -243,6 +256,7 @@ def track_glm_hierarchical_ess(self, step):
                 random_seed=100,
                 progressbar=False,
                 compute_convergence_checks=False,
+                return_inferencedata=False,
             )
             tot = time.time() - t0
         ess = float(az.ess(trace, var_names=["mu_a"])["mu_a"].values)
@@ -302,7 +316,9 @@ def freefall(y, t, p):
             Y = pm.Normal("Y", mu=ode_solution, sd=sigma, observed=y)
 
             t0 = time.time()
-            trace = pm.sample(500, tune=1000, chains=2, cores=2, random_seed=0)
+            trace = pm.sample(
+                500, tune=1000, chains=2, cores=2, random_seed=0, return_inferencedata=False
+            )
             tot = time.time() - t0
         ess = az.ess(trace)
         return np.mean([ess.sigma, ess.gamma]) / tot

pymc3/tests/sampler_fixtures.py

@@ -140,7 +140,13 @@ def setup_class(cls):
         cls.model = cls.make_model()
         with cls.model:
             cls.step = cls.make_step()
-            cls.trace = pm.sample(cls.n_samples, tune=cls.tune, step=cls.step, cores=cls.chains)
+            cls.trace = pm.sample(
+                cls.n_samples,
+                tune=cls.tune,
+                step=cls.step,
+                cores=cls.chains,
+                return_inferencedata=False,
+            )
         cls.samples = {}
         for var in cls.model.unobserved_RVs:
             cls.samples[get_var_name(var)] = cls.trace.get_values(var, burn=cls.burn)

pymc3/tests/test_bart.py

@@ -98,14 +98,14 @@ def test_model():
         sigma = pm.HalfNormal("sigma", 1)
         mu = pm.BART("mu", X, Y, m=50)
         y = pm.Normal("y", mu, sigma, observed=Y)
-        trace = pm.sample(1000, random_seed=212480)
+        trace = pm.sample(1000, random_seed=212480, return_inferencedata=False)
 
     np.testing.assert_allclose(trace[mu].mean(0), Y, 0.5)
 
     Y = np.repeat([0, 1], 50)
     with pm.Model() as model:
         mu = pm.BART("mu", X, Y, m=50, inv_link="logistic")
         y = pm.Bernoulli("y", mu, observed=Y)
-        trace = pm.sample(1000, random_seed=212480)
+        trace = pm.sample(1000, random_seed=212480, return_inferencedata=False)
 
     np.testing.assert_allclose(trace[mu].mean(0), Y, atol=0.5)

pymc3/tests/test_data_container.py

@@ -44,7 +44,7 @@ def test_sample(self):
             pm.Normal("obs", b * x_shared, np.sqrt(1e-2), observed=y)
 
             prior_trace0 = pm.sample_prior_predictive(1000)
-            trace = pm.sample(1000, init=None, tune=1000, chains=1)
+            trace = pm.sample(1000, init=None, tune=1000, chains=1, return_inferencedata=False)
             pp_trace0 = pm.sample_posterior_predictive(trace, 1000)
             pp_trace01 = pm.fast_sample_posterior_predictive(trace, 1000)
 
@@ -75,7 +75,7 @@ def test_sample_posterior_predictive_after_set_data(self):
             y = pm.Data("y", [1.0, 2.0, 3.0])
             beta = pm.Normal("beta", 0, 10.0)
             pm.Normal("obs", beta * x, np.sqrt(1e-2), observed=y)
-            trace = pm.sample(1000, tune=1000, chains=1)
+            trace = pm.sample(1000, tune=1000, chains=1, return_inferencedata=False)
         # Predict on new data.
         with model:
             x_test = [5.0, 6.0, 9.0]
@@ -94,13 +94,13 @@ def test_sample_after_set_data(self):
             y = pm.Data("y", [1.0, 2.0, 3.0])
             beta = pm.Normal("beta", 0, 10.0)
             pm.Normal("obs", beta * x, np.sqrt(1e-2), observed=y)
-            pm.sample(1000, init=None, tune=1000, chains=1)
+            pm.sample(1000, init=None, tune=1000, chains=1, return_inferencedata=False)
         # Predict on new data.
         new_x = [5.0, 6.0, 9.0]
         new_y = [5.0, 6.0, 9.0]
         with model:
             pm.set_data(new_data={"x": new_x, "y": new_y})
-            new_trace = pm.sample(1000, init=None, tune=1000, chains=1)
+            new_trace = pm.sample(1000, init=None, tune=1000, chains=1, return_inferencedata=False)
             pp_trace = pm.sample_posterior_predictive(new_trace, 1000)
             pp_tracef = pm.fast_sample_posterior_predictive(new_trace, 1000)
 
@@ -121,7 +121,7 @@ def test_shared_data_as_index(self):
             pm.Normal("obs", alpha[index], np.sqrt(1e-2), observed=y)
 
             prior_trace = pm.sample_prior_predictive(1000, var_names=["alpha"])
-            trace = pm.sample(1000, init=None, tune=1000, chains=1)
+            trace = pm.sample(1000, init=None, tune=1000, chains=1, return_inferencedata=False)
 
         # Predict on new data
         new_index = np.array([0, 1, 2])
@@ -146,14 +146,14 @@ def test_shared_data_as_rv_input(self):
         with pm.Model() as m:
             x = pm.Data("x", [1.0, 2.0, 3.0])
             _ = pm.Normal("y", mu=x, shape=3)
-            trace = pm.sample(chains=1)
+            trace = pm.sample(chains=1, return_inferencedata=False)
 
         np.testing.assert_allclose(np.array([1.0, 2.0, 3.0]), x.get_value(), atol=1e-1)
         np.testing.assert_allclose(np.array([1.0, 2.0, 3.0]), trace["y"].mean(0), atol=1e-1)
 
         with m:
             pm.set_data({"x": np.array([2.0, 4.0, 6.0])})
-            trace = pm.sample(chains=1)
+            trace = pm.sample(chains=1, return_inferencedata=False)
 
         np.testing.assert_allclose(np.array([2.0, 4.0, 6.0]), x.get_value(), atol=1e-1)
         np.testing.assert_allclose(np.array([2.0, 4.0, 6.0]), trace["y"].mean(0), atol=1e-1)
@@ -189,7 +189,7 @@ def test_set_data_to_non_data_container_variables(self):
             y = np.array([1.0, 2.0, 3.0])
             beta = pm.Normal("beta", 0, 10.0)
             pm.Normal("obs", beta * x, np.sqrt(1e-2), observed=y)
-            pm.sample(1000, init=None, tune=1000, chains=1)
+            pm.sample(1000, init=None, tune=1000, chains=1, return_inferencedata=False)
         with pytest.raises(TypeError) as error:
             pm.set_data({"beta": [1.1, 2.2, 3.3]}, model=model)
         error.match("defined as `pymc3.Data` inside the model")
@@ -201,7 +201,7 @@ def test_model_to_graphviz_for_model_with_data_container(self):
             beta = pm.Normal("beta", 0, 10.0)
             obs_sigma = floatX(np.sqrt(1e-2))
             pm.Normal("obs", beta * x, obs_sigma, observed=y)
-            pm.sample(1000, init=None, tune=1000, chains=1)
+            pm.sample(1000, init=None, tune=1000, chains=1, return_inferencedata=False)
 
         for formatting in {"latex", "latex_with_params"}:
             with pytest.raises(ValueError, match="Unsupported formatting"):

pymc3/tests/test_distributions.py

@@ -2596,7 +2596,7 @@ def func(x):
     with pm.Model():
         pm.Normal("x")
         y = pm.DensityDist("y", func)
-        pm.sample(draws=5, tune=1, mp_ctx="spawn")
+        pm.sample(draws=5, tune=1, mp_ctx="spawn", return_inferencedata=False)
 
     import pickle
 

pymc3/tests/test_distributions_random.py

@@ -1296,7 +1296,7 @@ def test_density_dist_with_random_sampleable(self, shape):
                 shape=shape,
                 random=normal_dist.random,
             )
-            trace = pm.sample(100, cores=1)
+            trace = pm.sample(100, cores=1, return_inferencedata=False)
 
         samples = 500
         size = 100
@@ -1319,7 +1319,7 @@ def test_density_dist_with_random_sampleable_failure(self, shape):
                 random=normal_dist.random,
                 wrap_random_with_dist_shape=False,
             )
-            trace = pm.sample(100, cores=1)
+            trace = pm.sample(100, cores=1, return_inferencedata=False)
 
         samples = 500
         with pytest.raises(RuntimeError):
@@ -1342,7 +1342,7 @@ def test_density_dist_with_random_sampleable_hidden_error(self, shape):
                 wrap_random_with_dist_shape=False,
                 check_shape_in_random=False,
             )
-            trace = pm.sample(100, cores=1)
+            trace = pm.sample(100, cores=1, return_inferencedata=False)
 
         samples = 500
         ppc = pm.sample_posterior_predictive(trace, samples=samples, model=model)
@@ -1365,7 +1365,7 @@ def test_density_dist_with_random_sampleable_handcrafted_success(self):
                 random=rvs,
                 wrap_random_with_dist_shape=False,
             )
-            trace = pm.sample(100, cores=1)
+            trace = pm.sample(100, cores=1, return_inferencedata=False)
 
         samples = 500
         size = 100
@@ -1385,7 +1385,7 @@ def test_density_dist_with_random_sampleable_handcrafted_success_fast(self):
                 random=rvs,
                 wrap_random_with_dist_shape=False,
             )
-            trace = pm.sample(100, cores=1)
+            trace = pm.sample(100, cores=1, return_inferencedata=False)
 
         samples = 500
         size = 100
@@ -1398,7 +1398,7 @@ def test_density_dist_without_random_not_sampleable(self):
             mu = pm.Normal("mu", 0, 1)
             normal_dist = pm.Normal.dist(mu, 1)
             pm.DensityDist("density_dist", normal_dist.logp, observed=np.random.randn(100))
-            trace = pm.sample(100, cores=1)
+            trace = pm.sample(100, cores=1, return_inferencedata=False)
 
         samples = 500
         with pytest.raises(ValueError):

pymc3/tests/test_examples.py

@@ -75,7 +75,7 @@ def build_model(self):
     def test_run(self):
         model = self.build_model()
         with model:
-            pm.sample(50, tune=50)
+            pm.sample(50, tune=50, return_inferencedata=False)
 
 
 class TestARM12_6(SeededTest):
@@ -112,7 +112,7 @@ def too_slow(self):
                 vars=[model["groupmean"], model["sd_interval__"], model["floor_m"]],
             )
             step = pm.NUTS(model.vars, scaling=start)
-            pm.sample(50, step=step, start=start)
+            pm.sample(50, step=step, start=start, return_inferencedata=False)
 
 
 class TestARM12_6Uranium(SeededTest):
@@ -159,7 +159,7 @@ def too_slow(self):
             h = np.diag(H(start))
 
             step = pm.HamiltonianMC(model.vars, h)
-            pm.sample(50, step=step, start=start)
+            pm.sample(50, step=step, start=start, return_inferencedata=False)
 
 
 def build_disaster_model(masked=False):
@@ -202,7 +202,9 @@ def test_disaster_model(self):
             start = {"early_mean": 2.0, "late_mean": 3.0}
             # Use slice sampler for means (other variables auto-selected)
             step = pm.Slice([model.early_mean_log__, model.late_mean_log__])
-            tr = pm.sample(500, tune=50, start=start, step=step, chains=2)
+            tr = pm.sample(
+                500, tune=50, start=start, step=step, chains=2, return_inferencedata=False
+            )
             az.summary(tr)
 
     def test_disaster_model_missing(self):
@@ -212,7 +214,9 @@ def test_disaster_model_missing(self):
             start = {"early_mean": 2.0, "late_mean": 3.0}
             # Use slice sampler for means (other variables auto-selected)
             step = pm.Slice([model.early_mean_log__, model.late_mean_log__])
-            tr = pm.sample(500, tune=50, start=start, step=step, chains=2)
+            tr = pm.sample(
+                500, tune=50, start=start, step=step, chains=2, return_inferencedata=False
+            )
             az.summary(tr)
 
 
@@ -231,7 +235,7 @@ def build_model(self):
     def test_run(self):
         with self.build_model():
             start = pm.find_MAP(method="Powell")
-            pm.sample(50, pm.Slice(), start=start)
+            pm.sample(50, pm.Slice(), start=start, return_inferencedata=False)
 
 
 class TestLatentOccupancy(SeededTest):
@@ -290,7 +294,9 @@ def test_run(self):
             }
             step_one = pm.Metropolis([model.theta_interval__, model.psi_logodds__])
             step_two = pm.BinaryMetropolis([model.z])
-            pm.sample(50, step=[step_one, step_two], start=start, chains=1)
+            pm.sample(
+                50, step=[step_one, step_two], start=start, chains=1, return_inferencedata=False
+            )
 
 
 @pytest.mark.xfail(
@@ -332,7 +338,7 @@ def build_model(self):
 
     def test_run(self):
         with self.build_model():
-            pm.sample(50, step=[pm.NUTS(), pm.Metropolis()])
+            pm.sample(50, step=[pm.NUTS(), pm.Metropolis()], return_inferencedata=False)
 
 
 class TestMultilevelNormal(SeededTest):
@@ -375,9 +381,9 @@ def test_run(self):
 
         with model:
             step = pm.MLDA(subsampling_rates=2, coarse_models=coarse_models)
-            pm.sample(draws=50, chains=2, tune=50, step=step)
+            pm.sample(draws=50, chains=2, tune=50, step=step, return_inferencedata=False)
 
             step = pm.MLDA(
                 subsampling_rates=2, coarse_models=coarse_models, base_sampler="Metropolis"
             )
-            pm.sample(draws=50, chains=2, tune=50, step=step)
+            pm.sample(draws=50, chains=2, tune=50, step=step, return_inferencedata=False)

pymc3/tests/test_hmc.py

@@ -53,7 +53,9 @@ def test_nuts_tuning():
     with model:
         pymc3.Normal("mu", mu=0, sigma=1)
         step = pymc3.NUTS()
-        trace = pymc3.sample(10, step=step, tune=5, progressbar=False, chains=1)
+        trace = pymc3.sample(
+            10, step=step, tune=5, progressbar=False, chains=1, return_inferencedata=False
+        )
 
     assert not step.tune
     assert np.all(trace["step_size"][5:] == trace["step_size"][5])

pymc3/tests/test_mixture.py

@@ -375,7 +375,7 @@ def build_toy_dataset(N, K):
 
             pm.Mixture("x_obs", pi, comp_dist, observed=X)
         with model:
-            trace = pm.sample(30, tune=10, chains=1)
+            trace = pm.sample(30, tune=10, chains=1, return_inferencedata=False)
 
         n_samples = 20
         with model:

pymc3/tests/test_model.py

@@ -194,7 +194,13 @@ def test_matrix_multiplication():
         rv_det = pm.Deterministic("rv_det", matrix @ rv_rv)
         det_rv = pm.Deterministic("det_rv", rv_rv @ transformed)
 
-        posterior = pm.sample(10, tune=0, compute_convergence_checks=False, progressbar=False)
+        posterior = pm.sample(
+            10,
+            tune=0,
+            compute_convergence_checks=False,
+            progressbar=False,
+            return_inferencedata=False,
+        )
         decimal = select_by_precision(7, 5)
         for point in posterior.points():
             npt.assert_almost_equal(