Vibe coded

Author: zaxtax

pymc/variational/approximations.py

@@ -85,6 +85,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:
@@ -139,6 +146,13 @@ 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")}
@@ -233,17 +247,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:

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"
             )