Make tests pass

Author: zaxtax

pymc/sampling/mcmc.py

@@ -1648,7 +1648,7 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
             compile_kwargs=compile_kwargs,
         )
         approx_sample = approx.sample(
-            draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
+            draws=chains, model=model, random_seed=random_seed_list[0], return_inferencedata=False
         )
         initial_points = [approx_sample[i] for i in range(chains)]
         std_apoint = approx.std.eval()
@@ -1672,7 +1672,7 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
             compile_kwargs=compile_kwargs,
         )
         approx_sample = approx.sample(
-            draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
+            draws=chains, model=model, random_seed=random_seed_list[0], return_inferencedata=False
         )
         initial_points = [approx_sample[i] for i in range(chains)]
         cov = approx.std.eval() ** 2
@@ -1690,7 +1690,7 @@ def model_logp_fn(ip: PointType) -> np.ndarray:
             compile_kwargs=compile_kwargs,
         )
         approx_sample = approx.sample(
-            draws=chains, random_seed=random_seed_list[0], return_inferencedata=False
+            draws=chains, model=model, random_seed=random_seed_list[0], return_inferencedata=False
         )
         initial_points = [approx_sample[i] for i in range(chains)]
         cov = approx.std.eval() ** 2

pymc/variational/operators.py

@@ -142,7 +142,7 @@ def __init__(self, approx, temperature=1):
 
     def apply(self, f):
         # f: kernel function for KSD f(histogram) -> (k(x,.), \nabla_x k(x,.))
-        if _known_scan_ignored_inputs([self.approx.model.logp()]):
+        if _known_scan_ignored_inputs([self.approx._model.logp()]):
             raise NotImplementedInference(
                 "SVGD does not currently support Minibatch or Simulator RV"
             )

pymc/variational/opvi.py

@@ -1289,7 +1289,7 @@ def symbolic_normalizing_constant(self):
                     obs.owner.inputs[1:],
                     constant_fold([obs.owner.inputs[0].shape], raise_not_constant=False),
                 )
-                for obs in self.model.observed_RVs
+                for obs in self._model.observed_RVs
                 if isinstance(obs.owner.op, MinibatchRandomVariable)
             ]
         )
@@ -1315,7 +1315,7 @@ def logq_norm(self):
     def _sized_symbolic_varlogp_and_datalogp(self):
         """*Dev* - computes sampled prior term from model via `pytensor.scan`."""
         varlogp_s, datalogp_s = self.symbolic_sample_over_posterior(
-            [self.model.varlogp, self.model.datalogp]
+            [self._model.varlogp, self._model.datalogp]
         )
         return varlogp_s, datalogp_s  # both shape (s,)
 
@@ -1352,7 +1352,7 @@ def datalogp(self):
     @node_property
     def _single_symbolic_varlogp_and_datalogp(self):
         """*Dev* - computes sampled prior term from model via `pytensor.scan`."""
-        varlogp, datalogp = self.symbolic_single_sample([self.model.varlogp, self.model.datalogp])
+        varlogp, datalogp = self.symbolic_single_sample([self._model.varlogp, self._model.datalogp])
         return varlogp, datalogp
 
     @node_property
@@ -1491,14 +1491,12 @@ def get_optimization_replacements(self, s, d):
         return repl
 
     @pytensor.config.change_flags(compute_test_value="off")
-    def sample_node(self, node, model=None, size=None, deterministic=False, more_replacements=None):
+    def sample_node(self, node, size=None, deterministic=False, more_replacements=None):
         """Sample given node or nodes over shared posterior.
 
         Parameters
         ----------
         node: PyTensor Variables (or PyTensor expressions)
-        model : Model (optional if in ``with`` context
-            Model to be used to generate samples.
         size: None or scalar
             number of samples
         more_replacements: `dict`
@@ -1513,7 +1511,7 @@ def sample_node(self, node, model=None, size=None, deterministic=False, more_rep
         """
         node_in = node
 
-        model = modelcontext(model)
+        model = self._model
 
         if more_replacements:
             node = graph_replace(node, more_replacements, strict=False)
@@ -1552,18 +1550,18 @@ def vars_names(vs):
         return found
 
     @node_property
-    def sample_dict_fn(self, model=None):
+    def sample_dict_fn(self):
         s = pt.iscalar()
 
-        model = modelcontext(model)
+        def inner(draws=100, *, model=None, random_seed: SeedSequenceSeed = None):
+            model = modelcontext(model)
 
-        names = [model.rvs_to_values[v].name for v in model.free_RVs]
-        sampled = [self.rslice(name, model) for name in names]
-        sampled = self.set_size_and_deterministic(sampled, s, 0)
-        sample_fn = compile([s], sampled)
-        rng_nodes = find_rng_nodes(sampled)
+            names = [model.rvs_to_values[v].name for v in model.free_RVs]
+            sampled = [self.rslice(name, model) for name in names]
+            sampled = self.set_size_and_deterministic(sampled, s, 0)
+            sample_fn = compile([s], sampled)
+            rng_nodes = find_rng_nodes(sampled)
 
-        def inner(draws=100, *, random_seed: SeedSequenceSeed = None):
             if random_seed is not None:
                 reseed_rngs(rng_nodes, random_seed)
             _samples = sample_fn(draws)

tests/variational/test_inference.py

@@ -41,7 +41,7 @@ def test_fit_with_nans(score):
         mean = inp * coef
         pm.Normal("y", mean, 0.1, observed=y)
         with pytest.raises(FloatingPointError) as e:
-            advi = pm.fit(100, score=score, obj_optimizer=pm.adam(learning_rate=float("nan")))
+            pm.fit(100, score=score, obj_optimizer=pm.adam(learning_rate=float("nan")))
 
 
 @pytest.fixture(scope="module", params=[True, False], ids=["mini", "full"])
@@ -174,8 +174,8 @@ 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(simple_model, inference, fit_kwargs, simple_model_data):
+    trace = inference.fit(**fit_kwargs).sample(10000, model=simple_model)
     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 +202,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))
 
@@ -269,7 +270,7 @@ def binomial_model_inference(binomial_model, inference_spec):
 def test_replacements(binomial_model_inference):
     d = pytensor.shared(1)
     approx = binomial_model_inference.approx
-    p = approx.model.p
+    p = approx._model.p
     p_t = p**3
     p_s = approx.sample_node(p_t)
     assert not any(
@@ -309,7 +310,7 @@ def test_sample_replacements(binomial_model_inference):
     i = pt.iscalar()
     i.tag.test_value = 1
     approx = binomial_model_inference.approx
-    p = approx.model.p
+    p = approx._model.p
     p_t = p**3
     p_s = approx.sample_node(p_t, size=100)
     if pytensor.config.compute_test_value != "off":

tests/variational/test_opvi.py

@@ -213,19 +213,19 @@ def test_pickle_approx(three_var_approx):
 
     dump = cloudpickle.dumps(three_var_approx)
     new = cloudpickle.loads(dump)
-    assert new.sample(1)
+    assert new.sample(1, model=new._model)
 
 
 def test_pickle_single_group(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)
+    assert new.sample(1, model=new._model)
 
 
 def test_sample_simple(three_var_approx):
-    trace = three_var_approx.sample(100, return_inferencedata=False)
+    trace = three_var_approx.sample(100, model=three_var_approx._model, return_inferencedata=False)
     assert set(trace.varnames) == {"one", "one_log__", "three", "two"}
     assert len(trace) == 100
     assert trace[0]["one"].shape == (10, 2)