Approximation's sample method uses model contexts

Author: zaxtax

pymc/sampling/mcmc.py

@@ -1731,9 +1731,10 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
             obj_optimizer=pm.adagrad_window,
             compile_kwargs=compile_kwargs,
         )
-        approx_sample = approx.sample(
-            draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
-        )
+        with model:
+            approx_sample = approx.sample(
+                draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
+            )
         initial_points = [
             {k: np.asarray(v) for k, v in approx_sample[i].items()} for i in range(chains)
         ]
@@ -1757,9 +1758,10 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
             obj_optimizer=pm.adagrad_window,
             compile_kwargs=compile_kwargs,
         )
-        approx_sample = approx.sample(
-            draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
-        )
+        with model:
+            approx_sample = approx.sample(
+                draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
+            )
         initial_points = [
             {k: np.asarray(v) for k, v in approx_sample[i].items()} for i in range(chains)
         ]
@@ -1777,9 +1779,10 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
             obj_optimizer=pm.adagrad_window,
             compile_kwargs=compile_kwargs,
         )
-        approx_sample = approx.sample(
-            draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
-        )
+        with model:
+            approx_sample = approx.sample(
+                draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
+            )
         initial_points = [
             {k: np.asarray(v) for k, v in approx_sample[i].items()} for i in range(chains)
         ]

pymc/variational/approximations.py

@@ -329,7 +329,8 @@ def sample_approx(approx, draws=100, include_transformed=True):
     trace: class:`pymc.backends.base.MultiTrace`
         Samples drawn from variational posterior.
     """
-    return approx.sample(draws=draws, include_transformed=include_transformed)
+    with approx.model:
+        return approx.sample(draws=draws, include_transformed=include_transformed)
 
 
 # single group shortcuts exported to user

pymc/variational/opvi.py

@@ -1229,7 +1229,10 @@ def __init__(self, groups, model=None):
             else:
                 rest.__init_group__(unseen_free_RVs)
                 self.groups.append(rest)
-        self.model = model
+
+    @property
+    def model(self):
+        return modelcontext(self.groups[0].model if self.groups else None)
 
     @property
     def has_logq(self):
@@ -1503,14 +1506,8 @@ def rslice(self, name):
 
         This node still needs :func:`set_size_and_deterministic` to be evaluated.
         """
-
-        def vars_names(vs):
-            return {self.model.rvs_to_values[v].name for v in vs}
-
-        for vars_, random, ordering in zip(
-            self.collect("group"), self.symbolic_randoms, self.collect("ordering")
-        ):
-            if name in vars_names(vars_):
+        for random, ordering in zip(self.symbolic_randoms, self.collect("ordering")):
+            if name in ordering:
                 name_, slc, shape, dtype = ordering[name]
                 found = random[..., slc].reshape((random.shape[0], *shape)).astype(dtype)
                 found.name = name + "_vi_random_slice"
@@ -1522,7 +1519,7 @@ def vars_names(vs):
     @node_property
     def sample_dict_fn(self):
         s = pt.iscalar()
-        names = [self.model.rvs_to_values[v].name for v in self.model.free_RVs]
+        names = [name for ordering in self.collect("ordering") for name in ordering]
         sampled = [self.rslice(name) for name in names]
         sampled = self.set_size_and_deterministic(sampled, s, 0)
         sample_fn = compile([s], sampled)
@@ -1567,8 +1564,10 @@ def sample(
             for i in range(draws)
         )
 
+        model = modelcontext(None)
+
         trace = NDArray(
-            model=self.model,
+            model=model,
             test_point={name: np.asarray(records[0]) for name, records in samples.items()},
         )
         try:
@@ -1582,7 +1581,7 @@ def sample(
         if not return_inferencedata:
             return multi_trace
         else:
-            return pm.to_inference_data(multi_trace, model=self.model, **kwargs)
+            return pm.to_inference_data(multi_trace, model=model, **kwargs)
 
     @property
     def ndim(self):

tests/variational/test_inference.py

@@ -14,6 +14,7 @@
 
 import io
 import operator
+import warnings
 
 import cloudpickle
 import numpy as np
@@ -24,6 +25,7 @@
 import pymc as pm
 import pymc.variational.opvi as opvi
 
+from pymc.model.transform.basic import remove_minibatched_nodes
 from pymc.variational.inference import ADVI, ASVGD, SVGD, FullRankADVI
 from pymc.variational.opvi import NotImplementedInference
 from tests import models
@@ -174,8 +176,9 @@ def fit_kwargs(inference, use_minibatch):
     return _select[(type(inference), key)]
 
 
-def test_fit_oo(inference, fit_kwargs, simple_model_data):
-    trace = inference.fit(**fit_kwargs).sample(10000)
+def test_fit_oo(inference, fit_kwargs, simple_model, simple_model_data):
+    with simple_model:
+        trace = inference.fit(**fit_kwargs).sample(10000)
     mu_post = simple_model_data["mu_post"]
     d = simple_model_data["d"]
     np.testing.assert_allclose(np.mean(trace.posterior["mu"]), mu_post, rtol=0.05)
@@ -202,7 +205,8 @@ def test_fit_start(inference_spec, simple_model):
         inference = inference_spec(**kw)
 
     try:
-        trace = inference.fit(n=0).sample(10000)
+        with simple_model:
+            trace = inference.fit(n=0).sample(10000)
     except NotImplementedInference as e:
         pytest.skip(str(e))
 
@@ -444,6 +448,40 @@ def test_fit_data_coords(hierarchical_model, hierarchical_model_data):
         assert data["mu"].shape == ()
 
 
+def test_sample_posterior_after_minibatch():
+    with pm.Model(coords={"obs_id": [0, 1, 2]}) as model:
+        x = pm.Data("x", [1.0, 2.0, 3.0], dims="obs_id")
+        y = pm.Data("y", [1.0, 2.0, 3.0], dims="obs_id")
+        x_mini, y_mini = pm.Minibatch(x, y, batch_size=2)
+        beta = pm.Normal("beta", 0, 10.0)
+        y_hat = pm.Deterministic("y_hat", beta * x_mini, dims="obs_id")
+        pm.Normal("obs", y_hat, np.sqrt(1e-2), observed=y_mini, total_size=3, dims="obs_id")
+        approx = pm.fit(
+            10,
+            method="advi",
+            progressbar=False,
+        )
+
+    model_post = remove_minibatched_nodes(model)
+    with model_post:
+        trace = approx.sample(500)
+
+    assert trace.posterior["beta"].shape == (1, 500)
+    assert trace.constant_data["x"].shape == (3,)
+    assert trace.observed_data["obs"].shape == (3,)
+
+    with model_post, warnings.catch_warnings():
+        warnings.filterwarnings("ignore", "Numba will use object mode", UserWarning)
+        x_test = [5, 6, 9, 12, 15]
+        pm.set_data(
+            new_data={"x": x_test, "y": [0.0] * len(x_test)},
+            coords={"obs_id": list(range(len(x_test)))},
+        )
+        y_test = pm.sample_posterior_predictive(trace, predictions=True, progressbar=False)
+
+    assert y_test.predictions["obs"].shape == (1, 500, 5)
+
+
 def test_multiple_minibatch_variables():
     """Regression test for bug reported in
     https://discourse.pymc.io/t/verifying-that-minibatch-is-actually-randomly-sampling/14308

tests/variational/test_opvi.py

@@ -208,24 +208,27 @@ def three_var_approx_single_group_mf(three_var_model):
     return MeanField(model=three_var_model)
 
 
-def test_pickle_approx(three_var_approx):
+def test_pickle_approx(three_var_model, three_var_approx):
     import cloudpickle
 
     dump = cloudpickle.dumps(three_var_approx)
     new = cloudpickle.loads(dump)
-    assert new.sample(1)
+    with three_var_model:
+        assert new.sample(1)
 
 
-def test_pickle_single_group(three_var_approx_single_group_mf):
+def test_pickle_single_group(three_var_model, three_var_approx_single_group_mf):
     import cloudpickle
 
     dump = cloudpickle.dumps(three_var_approx_single_group_mf)
     new = cloudpickle.loads(dump)
-    assert new.sample(1)
+    with three_var_model:
+        assert new.sample(1)
 
 
-def test_sample_simple(three_var_approx):
-    trace = three_var_approx.sample(100, return_inferencedata=False)
+def test_sample_simple(three_var_model, three_var_approx):
+    with three_var_model:
+        trace = three_var_approx.sample(100, return_inferencedata=False)
     assert set(trace.varnames) == {"one", "one_log__", "three", "two"}
     assert len(trace) == 100
     assert trace[0]["one"].shape == (10, 2)