Use better assert_no_rvs from logprob submodule

This utility can find RVs in inner graphs

Author: ricardoV94

pymc/testing.py

@@ -24,7 +24,8 @@
 from numpy import random as nr
 from numpy import testing as npt
 from pytensor.compile.mode import Mode
-from pytensor.graph.basic import ancestors
+from pytensor.graph.basic import walk
+from pytensor.graph.op import HasInnerGraph
 from pytensor.graph.rewriting.basic import in2out
 from pytensor.tensor import TensorVariable
 from pytensor.tensor.random.op import RandomVariable
@@ -37,7 +38,7 @@
 from pymc.distributions.shape_utils import change_dist_size
 from pymc.initial_point import make_initial_point_fn
 from pymc.logprob import joint_logp
-from pymc.logprob.abstract import icdf
+from pymc.logprob.abstract import MeasurableVariable, icdf
 from pymc.logprob.utils import ParameterValueError
 from pymc.pytensorf import (
     compile_pymc,
@@ -958,5 +959,18 @@ def seeded_numpy_distribution_builder(dist_name: str) -> Callable:
 
 
 def assert_no_rvs(var):
-    assert not any(isinstance(v.owner.op, RandomVariable) for v in ancestors([var]) if v.owner)
-    return var
+    """Assert that there are no `MeasurableVariable` nodes in a graph."""
+
+    def expand(r):
+        owner = r.owner
+        if owner:
+            inputs = list(reversed(owner.inputs))
+
+            if isinstance(owner.op, HasInnerGraph):
+                inputs += owner.op.inner_outputs
+
+            return inputs
+
+    for v in walk([var], expand, False):
+        if v.owner and isinstance(v.owner.op, (RandomVariable, MeasurableVariable)):
+            raise AssertionError(f"RV found in graph: {v}")

tests/logprob/utils.py

@@ -39,19 +39,11 @@
 import numpy as np
 
 from pytensor import tensor as pt
-from pytensor.graph.basic import walk
-from pytensor.graph.op import HasInnerGraph
 from pytensor.tensor.var import TensorVariable
 from scipy import stats as stats
 
 from pymc.logprob import factorized_joint_logprob
-from pymc.logprob.abstract import (
-    MeasurableVariable,
-    get_measurable_outputs,
-    icdf,
-    logcdf,
-    logprob,
-)
+from pymc.logprob.abstract import get_measurable_outputs, icdf, logcdf, logprob
 from pymc.logprob.utils import ignore_logprob
 
 
@@ -82,24 +74,6 @@ def joint_logprob(*args, sum: bool = True, **kwargs) -> Optional[TensorVariable]
     return pt.add(*logprob.values())
 
 
-def assert_no_rvs(var):
-    """Assert that there are no `MeasurableVariable` nodes in a graph."""
-
-    def expand(r):
-        owner = r.owner
-        if owner:
-            inputs = list(reversed(owner.inputs))
-
-            if isinstance(owner.op, HasInnerGraph):
-                inputs += owner.op.inner_outputs
-
-            return inputs
-
-    for v in walk([var], expand, False):
-        if v.owner and isinstance(v.owner.op, MeasurableVariable):
-            raise AssertionError(f"Variable {v} is a MeasurableVariable")
-
-
 def simulate_poiszero_hmm(
     N, mu=10.0, pi_0_a=np.r_[1, 1], p_0_a=np.r_[5, 1], p_1_a=np.r_[1, 1], seed=None
 ):