Cleanup RV in graph checks

Author: ricardoV94

pymc/logprob/basic.py

@@ -40,7 +40,6 @@
 from typing import Dict, List, Optional, Sequence, Union
 
 import numpy as np
-import pytensor
 import pytensor.tensor as pt
 
 from pytensor import config
@@ -53,7 +52,6 @@
 )
 from pytensor.graph.op import compute_test_value
 from pytensor.graph.rewriting.basic import GraphRewriter, NodeRewriter
-from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
@@ -66,11 +64,22 @@
 )
 from pymc.logprob.rewriting import cleanup_ir, construct_ir_fgraph
 from pymc.logprob.transforms import RVTransform, TransformValuesRewrite
-from pymc.logprob.utils import rvs_to_value_vars
+from pymc.logprob.utils import find_rvs_in_graph, rvs_to_value_vars
 
 TensorLike: TypeAlias = Union[Variable, float, np.ndarray]
 
 
+def _find_unallowed_rvs_in_graph(graph):
+    from pymc.data import MinibatchIndexRV
+    from pymc.distributions.simulator import SimulatorRV
+
+    return {
+        rv
+        for rv in find_rvs_in_graph(graph)
+        if not isinstance(rv.owner.op, (SimulatorRV, MinibatchIndexRV))
+    }
+
+
 def _warn_rvs_in_inferred_graph(graph: Union[TensorVariable, Sequence[TensorVariable]]):
     """Issue warning if any RVs are found in graph.
 
@@ -81,13 +90,11 @@ def _warn_rvs_in_inferred_graph(graph: Union[TensorVariable, Sequence[TensorVari
     This makes it impossible (or difficult) to replace it by the respective values afterward,
     so we instruct users to do it beforehand.
     """
-    from pymc.testing import assert_no_rvs
 
-    try:
-        assert_no_rvs(graph)
-    except AssertionError:
+    rvs_in_graph = _find_unallowed_rvs_in_graph(graph)
+    if rvs_in_graph:
         warnings.warn(
-            "RandomVariables were found in the derived graph. "
+            f"RandomVariables {rvs_in_graph} were found in the derived graph. "
             "These variables are a clone and do not match the original ones on identity.\n"
             "If you are deriving a quantity that depends on model RVs, use `model.replace_rvs_by_values` first. For example: "
             "`logp(model.replace_rvs_by_values([rv])[0], value)`",
@@ -149,6 +156,13 @@ def icdf(
         return expr
 
 
+RVS_IN_JOINT_LOGP_GRAPH_MSG = (
+    "Random variables detected in the logp graph: %s.\n"
+    "This can happen when DensityDist logp or Interval transform functions reference nonlocal variables,\n"
+    "or when not all rvs have a corresponding value variable."
+)
+
+
 def factorized_joint_logprob(
     rv_values: Dict[TensorVariable, TensorVariable],
     warn_missing_rvs: bool = True,
@@ -316,34 +330,13 @@ def factorized_joint_logprob(
     cleanup_ir(logprob_expressions)
 
     if warn_missing_rvs:
-        _warn_rvs_in_inferred_graph(logprob_expressions)
+        rvs_in_logp_expressions = _find_unallowed_rvs_in_graph(logprob_expressions)
+        if rvs_in_logp_expressions:
+            warnings.warn(RVS_IN_JOINT_LOGP_GRAPH_MSG % rvs_in_logp_expressions, UserWarning)
 
     return logprob_vars
 
 
-def _check_no_rvs(logp_terms: Sequence[TensorVariable]):
-    # Raise if there are unexpected RandomVariables in the logp graph
-    # Only SimulatorRVs MinibatchIndexRVs are allowed
-    from pymc.data import MinibatchIndexRV
-    from pymc.distributions.simulator import SimulatorRV
-
-    unexpected_rv_nodes = [
-        node
-        for node in pytensor.graph.ancestors(logp_terms)
-        if (
-            node.owner
-            and isinstance(node.owner.op, RandomVariable)
-            and not isinstance(node.owner.op, (SimulatorRV, MinibatchIndexRV))
-        )
-    ]
-    if unexpected_rv_nodes:
-        raise ValueError(
-            f"Random variables detected in the logp graph: {unexpected_rv_nodes}.\n"
-            "This can happen when DensityDist logp or Interval transform functions "
-            "reference nonlocal variables."
-        )
-
-
 def joint_logp(
     rvs: Sequence[TensorVariable],
     *,
@@ -381,5 +374,10 @@ def joint_logp(
         value_var = rvs_to_values[rv]
         logp_terms[value_var] = temp_logp_terms[value_var]
 
-    _check_no_rvs(list(logp_terms.values()))
-    return list(logp_terms.values())
+    logp_terms_list = list(logp_terms.values())
+
+    rvs_in_logp_expressions = _find_unallowed_rvs_in_graph(logp_terms_list)
+    if rvs_in_logp_expressions:
+        raise ValueError(RVS_IN_JOINT_LOGP_GRAPH_MSG % rvs_in_logp_expressions)
+
+    return logp_terms_list

pymc/logprob/utils.py

@@ -36,19 +36,34 @@
 
 import warnings
 
-from typing import Callable, Dict, Generator, Iterable, List, Optional, Set, Tuple
+from typing import (
+    Callable,
+    Dict,
+    Generator,
+    Iterable,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 
+from pytensor import Variable
 from pytensor import tensor as pt
 from pytensor.graph import Apply, Op
 from pytensor.graph.basic import Constant, clone_get_equiv, graph_inputs, walk
 from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.op import HasInnerGraph
 from pytensor.link.c.type import CType
 from pytensor.raise_op import CheckAndRaise
+from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.var import TensorVariable
 
 from pymc.logprob.abstract import MeasurableVariable, _logprob
+from pymc.util import makeiter
 
 
 def walk_model(
@@ -273,3 +288,23 @@ def diracdelta_logprob(op, values, *inputs, **kwargs):
     (const_value,) = inputs
     values, const_value = pt.broadcast_arrays(values, const_value)
     return pt.switch(pt.isclose(values, const_value, rtol=op.rtol, atol=op.atol), 0.0, -np.inf)
+
+
+def find_rvs_in_graph(vars: Union[Variable, Sequence[Variable]]) -> Set[Variable]:
+    """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
+
+    return {
+        node
+        for node in walk(makeiter(vars), expand, False)
+        if node.owner and isinstance(node.owner.op, (RandomVariable, MeasurableVariable))
+    }

pymc/pytensorf.py

@@ -65,6 +65,7 @@
 from pymc.exceptions import NotConstantValueError
 from pymc.logprob.transforms import RVTransform
 from pymc.logprob.utils import CheckParameterValue
+from pymc.util import makeiter
 from pymc.vartypes import continuous_types, isgenerator, typefilter
 
 PotentialShapeType = Union[int, np.ndarray, Sequence[Union[int, Variable]], TensorVariable]
@@ -550,13 +551,6 @@ def hessian_diag(f, vars=None):
         return empty_gradient
 
 
-def makeiter(a):
-    if isinstance(a, (tuple, list)):
-        return a
-    else:
-        return [a]
-
-
 class IdentityOp(scalar.UnaryScalarOp):
     @staticmethod
     def st_impl(x):

pymc/testing.py

@@ -24,11 +24,9 @@
 from numpy import random as nr
 from numpy import testing as npt
 from pytensor.compile.mode import Mode
-from pytensor.graph.basic import Variable, walk
-from pytensor.graph.op import HasInnerGraph
+from pytensor.graph.basic import Variable
 from pytensor.graph.rewriting.basic import in2out
 from pytensor.tensor import TensorVariable
-from pytensor.tensor.random.op import RandomVariable
 from scipy import special as sp
 from scipy import stats as st
 
@@ -37,16 +35,14 @@
 from pymc.distributions.distribution import Distribution
 from pymc.distributions.shape_utils import change_dist_size
 from pymc.initial_point import make_initial_point_fn
-from pymc.logprob.abstract import MeasurableVariable
 from pymc.logprob.basic import icdf, joint_logp, logcdf, logp
-from pymc.logprob.utils import ParameterValueError
+from pymc.logprob.utils import ParameterValueError, find_rvs_in_graph
 from pymc.pytensorf import (
     compile_pymc,
     floatX,
     inputvars,
     intX,
     local_check_parameter_to_ninf_switch,
-    makeiter,
 )
 
 # This mode can be used for tests where model compilations takes the bulk of the runtime
@@ -964,19 +960,9 @@ def seeded_numpy_distribution_builder(dist_name: str) -> Callable:
     )
 
 
-def assert_no_rvs(vars: Union[Variable, Sequence[Variable]]):
+def assert_no_rvs(vars: Sequence[Variable]) -> None:
     """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(makeiter(vars), expand, False):
-        if v.owner and isinstance(v.owner.op, (RandomVariable, MeasurableVariable)):
-            raise AssertionError(f"RV found in graph: {v}")
+    rvs = find_rvs_in_graph(vars)
+    if rvs:
+        raise AssertionError(f"RV found in graph: {rvs}")

pymc/util.py

@@ -510,3 +510,10 @@ def _add_future_warning_tag(var) -> None:
         for k, v in old_tag.__dict__.items():
             new_tag.__dict__.setdefault(k, v)
         var.tag = new_tag
+
+
+def makeiter(a):
+    if isinstance(a, (tuple, list)):
+        return a
+    else:
+        return [a]

pymc/variational/approximations.py

@@ -25,7 +25,7 @@
 
 from pymc.blocking import DictToArrayBijection
 from pymc.distributions.dist_math import rho2sigma
-from pymc.pytensorf import makeiter
+from pymc.util import makeiter
 from pymc.variational import opvi
 from pymc.variational.opvi import (
     Approximation,

pymc/variational/opvi.py

@@ -72,14 +72,14 @@
     compile_pymc,
     find_rng_nodes,
     identity,
-    makeiter,
     reseed_rngs,
 )
 from pymc.util import (
     RandomState,
     WithMemoization,
     _get_seeds_per_chain,
     locally_cachedmethod,
+    makeiter,
 )
 from pymc.variational.minibatch_rv import MinibatchRandomVariable, get_scaling
 from pymc.variational.updates import adagrad_window

tests/logprob/test_basic.py

@@ -223,7 +223,7 @@ def test_warn_random_found_factorized_joint_logprob():
 
     y_vv = y_rv.clone()
 
-    with pytest.warns(UserWarning, match="RandomVariables were found in the derived graph"):
+    with pytest.warns(UserWarning, match="Random variables detected in the logp graph: {x}"):
         factorized_joint_logprob({y_rv: y_vv})
 
     with warnings.catch_warnings():
@@ -443,7 +443,9 @@ def test_warn_random_found_probability_inference(func, scipy_func, test_value):
         # In which case the inference should either return that or fail explicitly
         # For now, the lopgrob submodule treats the input as a stochastic value.
         rv = pt.exp(pm.Normal.dist(input_rv))
-        with pytest.warns(UserWarning, match="RandomVariables were found in the derived graph"):
+        with pytest.warns(
+            UserWarning, match="RandomVariables {input} were found in the derived graph"
+        ):
             assert func(rv, 0.0)
 
         res = func(rv, 0.0, warn_missing_rvs=False)