Merge functionality of pytensorf and logprob/utils

Also fixes circular imports

Author: ricardoV94

pymc/gp/cov.py

@@ -48,6 +48,8 @@
     "Kron",
 ]
 
+from pymc.pytensorf import constant_fold
+
 TensorLike = Union[np.ndarray, TensorVariable]
 IntSequence = Union[np.ndarray, Sequence[int]]
 
@@ -183,9 +185,6 @@ def n_dims(self) -> int:
     def _slice(self, X, Xs=None):
         xdims = X.shape[-1]
         if isinstance(xdims, Variable):
-            # Circular dependency
-            from pymc.pytensorf import constant_fold
-
             [xdims] = constant_fold([xdims])
         if self.input_dim != xdims:
             warnings.warn(

pymc/gp/util.py

@@ -18,13 +18,14 @@
 import pytensor.tensor as pt
 
 from pytensor.compile import SharedVariable
+from pytensor.graph import ancestors
 from pytensor.tensor.variable import TensorConstant
 from scipy.cluster.vq import kmeans
 
 # Avoid circular dependency when importing modelcontext
 from pymc.distributions.distribution import Distribution
 from pymc.model import modelcontext
-from pymc.pytensorf import compile_pymc, walk_model
+from pymc.pytensorf import compile_pymc
 
 _ = Distribution  # keep both pylint and black happy
 
@@ -48,7 +49,7 @@ def replace_with_values(vars_needed, replacements=None, model=None):
     model = modelcontext(model)
 
     inputs, input_names = [], []
-    for rv in walk_model(vars_needed):
+    for rv in ancestors(vars_needed):
         if rv in model.named_vars.values() and not isinstance(rv, SharedVariable):
             inputs.append(rv)
             input_names.append(rv.name)

pymc/logprob/basic.py

@@ -65,7 +65,8 @@
 from pymc.logprob.rewriting import cleanup_ir, construct_ir_fgraph
 from pymc.logprob.transform_value import TransformValuesRewrite
 from pymc.logprob.transforms import Transform
-from pymc.logprob.utils import find_rvs_in_graph, rvs_to_value_vars
+from pymc.logprob.utils import rvs_in_graph
+from pymc.pytensorf import replace_vars_in_graphs
 
 TensorLike: TypeAlias = Union[Variable, float, np.ndarray]
 
@@ -76,7 +77,7 @@ def _find_unallowed_rvs_in_graph(graph):
 
     return {
         rv
-        for rv in find_rvs_in_graph(graph)
+        for rv in rvs_in_graph(graph)
         if not isinstance(rv.owner.op, (SimulatorRV, MinibatchIndexRV))
     }
 
@@ -530,11 +531,9 @@ def conditional_logp(
             continue
 
         # Replace `RandomVariable`s in the inputs with value variables.
-        # Also, store the results in the `replacements` map for the nodes
-        # that follow.
-        remapped_vars, _ = rvs_to_value_vars(
-            q_values + list(node.inputs),
-            initial_replacements=replacements,
+        remapped_vars = replace_vars_in_graphs(
+            graphs=q_values + list(node.inputs),
+            replacements=replacements,
         )
         q_values = remapped_vars[: len(q_values)]
         q_rv_inputs = remapped_vars[len(q_values) :]
@@ -562,8 +561,7 @@ def conditional_logp(
 
             logprob_vars[q_value_var] = q_logprob_var
 
-        # Recompute test values for the changes introduced by the
-        # replacements above.
+        # Recompute test values for the changes introduced by the replacements above.
         if config.compute_test_value != "off":
             for node in io_toposort(graph_inputs(q_logprob_vars), q_logprob_vars):
                 compute_test_value(node)

pymc/logprob/checks.py

@@ -44,6 +44,7 @@
 
 from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
+from pymc.logprob.utils import replace_rvs_by_values
 
 
 class MeasurableSpecifyShape(SpecifyShape):
@@ -107,8 +108,6 @@ class MeasurableCheckAndRaise(CheckAndRaise):
 
 @_logprob.register(MeasurableCheckAndRaise)
 def logprob_check_and_raise(op, values, inner_rv, *assertions, **kwargs):
-    from pymc.pytensorf import replace_rvs_by_values
-
     (value,) = values
     # transfer assertion from rv to value
     assertions = replace_rvs_by_values(assertions, rvs_to_values={inner_rv: value})

pymc/logprob/mixture.py

@@ -78,7 +78,8 @@
     measurable_ir_rewrites_db,
     subtensor_ops,
 )
-from pymc.logprob.utils import check_potential_measurability
+from pymc.logprob.utils import check_potential_measurability, replace_rvs_by_values
+from pymc.pytensorf import constant_fold
 
 
 def is_newaxis(x):
@@ -255,9 +256,6 @@ def get_stack_mixture_vars(
         mixture_rvs = joined_rvs.owner.inputs
 
     elif isinstance(joined_rvs.owner.op, Join):
-        # TODO: Find better solution to avoid this circular dependency
-        from pymc.pytensorf import constant_fold
-
         join_axis = joined_rvs.owner.inputs[0]
         # TODO: Support symbolic join axes. This will raise ValueError if it's not a constant
         (join_axis,) = constant_fold((join_axis,), raise_not_constant=False)
@@ -351,9 +349,6 @@ def logprob_MixtureRV(
             comp_rvs = [comp[None] for comp in comp_rvs]
             original_shape = (len(comp_rvs),)
         else:
-            # TODO: Find better solution to avoid this circular dependency
-            from pymc.pytensorf import constant_fold
-
             join_axis_val = constant_fold((join_axis,))[0].item()
             original_shape = shape_tuple(comp_rvs[0])
 
@@ -544,7 +539,6 @@ def find_measurable_ifelse_mixture(fgraph, node):
 @_logprob.register(MeasurableIfElse)
 def logprob_ifelse(op, values, if_var, *base_rvs, **kwargs):
     """Compute the log-likelihood graph for an `IfElse`."""
-    from pymc.pytensorf import replace_rvs_by_values
 
     assert len(values) * 2 == len(base_rvs)
 

pymc/logprob/scan.py

@@ -62,7 +62,7 @@
     logprob_rewrites_db,
     measurable_ir_rewrites_db,
 )
-from pymc.pytensorf import replace_rvs_by_values
+from pymc.logprob.utils import replace_rvs_by_values
 
 
 class MeasurableScan(Scan):

pymc/logprob/tensor.py

@@ -55,7 +55,8 @@
     assume_measured_ir_outputs,
     measurable_ir_rewrites_db,
 )
-from pymc.logprob.utils import check_potential_measurability
+from pymc.logprob.utils import check_potential_measurability, replace_rvs_by_values
+from pymc.pytensorf import constant_fold
 
 
 @node_rewriter([Alloc])
@@ -131,7 +132,6 @@ class MeasurableMakeVector(MakeVector):
 def logprob_make_vector(op, values, *base_rvs, **kwargs):
     """Compute the log-likelihood graph for a `MeasurableMakeVector`."""
     # TODO: Sort out this circular dependency issue
-    from pymc.pytensorf import replace_rvs_by_values
 
     (value,) = values
 
@@ -158,9 +158,6 @@ class MeasurableJoin(Join):
 @_logprob.register(MeasurableJoin)
 def logprob_join(op, values, axis, *base_rvs, **kwargs):
     """Compute the log-likelihood graph for a `Join`."""
-    # TODO: Find better way to avoid circular dependency
-    from pymc.pytensorf import constant_fold, replace_rvs_by_values
-
     (value,) = values
 
     base_rv_shapes = [base_var.shape[axis] for base_var in base_rvs]