Approximation's sample method uses model contexts

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/approximations.py

@@ -12,6 +12,8 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
+from functools import cached_property
+
 import numpy as np
 import pytensor
 
@@ -32,7 +34,6 @@
     Group,
     NotImplementedInference,
     _known_scan_ignored_inputs,
-    node_property,
 )
 
 __all__ = ["Empirical", "FullRank", "MeanField", "sample_approx"]
@@ -52,20 +53,20 @@ class MeanFieldGroup(Group):
     short_name = "mean_field"
     alias_names = frozenset(["mf"])
 
-    @node_property
+    @cached_property
     def mean(self):
         return self.params_dict["mu"]
 
-    @node_property
+    @cached_property
     def rho(self):
         return self.params_dict["rho"]
 
-    @node_property
+    @cached_property
     def cov(self):
         var = rho2sigma(self.rho) ** 2
         return pt.diag(var)
 
-    @node_property
+    @cached_property
     def std(self):
         return rho2sigma(self.rho)
 
@@ -85,6 +86,13 @@ def create_shared_params(self, start=None, start_sigma=None):
         # by `self.ordering`. In the cases I looked into these turn out to be the same, but there may be edge cases or
         # future code changes that break this assumption.
         start = self._prepare_start(start)
+        # Ensure start is a 1D array and matches ddim
+        start = np.asarray(start).flatten()
+        if start.size != self.ddim:
+            raise ValueError(
+                f"Start array size mismatch: got {start.size}, expected {self.ddim}. "
+                f"Start shape: {start.shape if hasattr(start, 'shape') else 'unknown'}"
+            )
         rho1 = np.zeros((self.ddim,))
 
         if start_sigma is not None:
@@ -99,14 +107,14 @@ def create_shared_params(self, start=None, start_sigma=None):
             "rho": pytensor.shared(pm.floatX(rho), "rho"),
         }
 
-    @node_property
+    @cached_property
     def symbolic_random(self):
         initial = self.symbolic_initial
         sigma = self.std
         mu = self.mean
         return sigma * initial + mu
 
-    @node_property
+    @cached_property
     def symbolic_logq_not_scaled(self):
         z0 = self.symbolic_initial
         std = rho2sigma(self.rho)
@@ -139,28 +147,35 @@ def __init_group__(self, group):
 
     def create_shared_params(self, start=None):
         start = self._prepare_start(start)
+        # Ensure start is a 1D array and matches ddim
+        start = np.asarray(start).flatten()
+        if start.size != self.ddim:
+            raise ValueError(
+                f"Start array size mismatch: got {start.size}, expected {self.ddim}. "
+                f"Start shape: {start.shape if hasattr(start, 'shape') else 'unknown'}"
+            )
         n = self.ddim
         L_tril = np.eye(n)[np.tril_indices(n)].astype(pytensor.config.floatX)
         return {"mu": pytensor.shared(start, "mu"), "L_tril": pytensor.shared(L_tril, "L_tril")}
 
-    @node_property
+    @cached_property
     def L(self):
         L = pt.zeros((self.ddim, self.ddim))
         L = pt.set_subtensor(L[self.tril_indices], self.params_dict["L_tril"])
         Ld = L[..., np.arange(self.ddim), np.arange(self.ddim)]
         L = pt.set_subtensor(Ld, rho2sigma(Ld))
         return L
 
-    @node_property
+    @cached_property
     def mean(self):
         return self.params_dict["mu"]
 
-    @node_property
+    @cached_property
     def cov(self):
         L = self.L
         return L.dot(L.T)
 
-    @node_property
+    @cached_property
     def std(self):
         return pt.sqrt(pt.diag(self.cov))
 
@@ -173,7 +188,7 @@ def num_tril_entries(self):
     def tril_indices(self):
         return np.tril_indices(self.ddim)
 
-    @node_property
+    @cached_property
     def symbolic_logq_not_scaled(self):
         z0 = self.symbolic_initial
         diag = pt.diagonal(self.L, 0, self.L.ndim - 2, self.L.ndim - 1)
@@ -182,7 +197,7 @@ def symbolic_logq_not_scaled(self):
         logq = quaddist - logdet
         return logq.sum(range(1, logq.ndim))
 
-    @node_property
+    @cached_property
     def symbolic_random(self):
         initial = self.symbolic_initial
         L = self.L
@@ -233,17 +248,19 @@ def create_shared_params(self, trace=None, size=None, jitter=1, start=None):
         return {"histogram": pytensor.shared(pm.floatX(histogram), "histogram")}
 
     def _check_trace(self):
+        from pymc.model import modelcontext
+
         trace = self._kwargs.get("trace", None)
         if isinstance(trace, InferenceData):
             raise NotImplementedError(
                 "The `Empirical` approximation does not yet support `InferenceData` inputs."
                 " Pass `pm.sample(return_inferencedata=False)` to get a `MultiTrace` to use with `Empirical`."
                 " Please help us to refactor: https://github.com/pymc-devs/pymc/issues/5884"
             )
-        elif trace is not None and not all(
-            self.model.rvs_to_values[var].name in trace.varnames for var in self.group
-        ):
-            raise ValueError("trace has not all free RVs in the group")
+        elif trace is not None:
+            model = modelcontext(None)
+            if not all(model.rvs_to_values[var].name in trace.varnames for var in self.group):
+                raise ValueError("trace has not all free RVs in the group")
 
     def randidx(self, size=None):
         if size is None:
@@ -284,24 +301,24 @@ def _new_initial(self, size, deterministic, more_replacements=None):
             else:
                 return self.histogram[self.randidx(size)]
 
-    @property
+    @cached_property
     def symbolic_random(self):
         return self.symbolic_initial
 
-    @property
+    @cached_property
     def histogram(self):
         return self.params_dict["histogram"]
 
-    @node_property
+    @cached_property
     def mean(self):
         return self.histogram.mean(0)
 
-    @node_property
+    @cached_property
     def cov(self):
         x = self.histogram - self.mean
         return x.T.dot(x) / pm.floatX(self.histogram.shape[0])
 
-    @node_property
+    @cached_property
     def std(self):
         return pt.sqrt(pt.diag(self.cov))
 

pymc/variational/operators.py

@@ -19,6 +19,7 @@
 
 import pymc as pm
 
+from pymc.model import modelcontext
 from pymc.variational import opvi
 from pymc.variational.opvi import (
     NotImplementedInference,
@@ -142,7 +143,8 @@ 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()]):
+        model = modelcontext(None)
+        if _known_scan_ignored_inputs([model.logp()]):
             raise NotImplementedInference(
                 "SVGD does not currently support Minibatch or Simulator RV"
             )