Use more readable ignore_logprob helper in logprob submodule

Author: ricardoV94

pymc/logprob/censoring.py

@@ -52,10 +52,9 @@
     MeasurableVariable,
     _logcdf,
     _logprob,
-    assign_custom_measurable_outputs,
 )
 from pymc.logprob.rewriting import measurable_ir_rewrites_db
-from pymc.logprob.utils import CheckParameterValue
+from pymc.logprob.utils import CheckParameterValue, ignore_logprob
 
 
 class MeasurableClip(MeasurableElemwise):
@@ -95,7 +94,7 @@ def find_measurable_clips(fgraph: FunctionGraph, node: Node) -> Optional[List[Me
     upper_bound = upper_bound if (upper_bound is not base_var) else pt.constant(np.inf)
 
     # Make base_var unmeasurable
-    unmeasurable_base_var = assign_custom_measurable_outputs(base_var.owner)
+    unmeasurable_base_var = ignore_logprob(base_var)
     clipped_rv_node = measurable_clip.make_node(unmeasurable_base_var, lower_bound, upper_bound)
     clipped_rv = clipped_rv_node.outputs[0]
 
@@ -198,7 +197,7 @@ def find_measurable_roundings(fgraph: FunctionGraph, node: Node) -> Optional[Lis
         return None
 
     # Make base_var unmeasurable
-    unmeasurable_base_var = assign_custom_measurable_outputs(base_var.owner)
+    unmeasurable_base_var = ignore_logprob(base_var)
 
     rounded_op = MeasurableRound(node.op.scalar_op)
     rounded_rv = rounded_op.make_node(unmeasurable_base_var).default_output()

pymc/logprob/cumsum.py

@@ -41,13 +41,9 @@
 from pytensor.graph.rewriting.basic import node_rewriter
 from pytensor.tensor.extra_ops import CumOp
 
-from pymc.logprob.abstract import (
-    MeasurableVariable,
-    _logprob,
-    assign_custom_measurable_outputs,
-    logprob,
-)
+from pymc.logprob.abstract import MeasurableVariable, _logprob, logprob
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
+from pymc.logprob.utils import ignore_logprob
 
 
 class MeasurableCumsum(CumOp):
@@ -112,7 +108,7 @@ def find_measurable_cumsums(fgraph, node) -> Optional[List[MeasurableCumsum]]:
 
     new_op = MeasurableCumsum(axis=node.op.axis or 0, mode="add")
     # Make base_var unmeasurable
-    unmeasurable_base_rv = assign_custom_measurable_outputs(base_rv.owner)
+    unmeasurable_base_rv = ignore_logprob(base_rv)
     new_rv = new_op.make_node(unmeasurable_base_rv).default_output()
     new_rv.name = rv.name
 

pymc/logprob/mixture.py

@@ -69,18 +69,14 @@
 from pytensor.tensor.type_other import NoneConst, NoneTypeT, SliceConstant, SliceType
 from pytensor.tensor.var import TensorVariable
 
-from pymc.logprob.abstract import (
-    MeasurableVariable,
-    _logprob,
-    assign_custom_measurable_outputs,
-    logprob,
-)
+from pymc.logprob.abstract import MeasurableVariable, _logprob, logprob
 from pymc.logprob.rewriting import (
     local_lift_DiracDelta,
     logprob_rewrites_db,
     subtensor_ops,
 )
 from pymc.logprob.tensor import naive_bcast_rv_lift
+from pymc.logprob.utils import ignore_logprob
 
 
 def is_newaxis(x):
@@ -328,9 +324,7 @@ def mixture_replace(fgraph, node):
         # We create custom types for the mixture components and assign them
         # null `get_measurable_outputs` dispatches so that they aren't
         # erroneously encountered in places like `factorized_joint_logprob`.
-        new_node = assign_custom_measurable_outputs(component_rv.owner)
-        out_idx = component_rv.owner.outputs.index(component_rv)
-        new_comp_rv = new_node.outputs[out_idx]
+        new_comp_rv = ignore_logprob(component_rv)
         new_mixture_rvs.append(new_comp_rv)
 
     # Replace this sub-graph with a `MixtureRV`
@@ -379,9 +373,7 @@ def switch_mixture_replace(fgraph, node):
             and component_rv not in rv_map_feature.rv_values
         ):
             return None
-        new_node = assign_custom_measurable_outputs(component_rv.owner)
-        out_idx = component_rv.owner.outputs.index(component_rv)
-        new_comp_rv = new_node.outputs[out_idx]
+        new_comp_rv = ignore_logprob(component_rv)
         mixture_rvs.append(new_comp_rv)
 
     mix_op = MixtureRV(

pymc/logprob/tensor.py

@@ -50,13 +50,9 @@
     local_rv_size_lift,
 )
 
-from pymc.logprob.abstract import (
-    MeasurableVariable,
-    _logprob,
-    assign_custom_measurable_outputs,
-    logprob,
-)
+from pymc.logprob.abstract import MeasurableVariable, _logprob, logprob
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
+from pymc.logprob.utils import ignore_logprob
 
 
 @node_rewriter([BroadcastTo])
@@ -233,12 +229,7 @@ def find_measurable_stacks(
         return None  # pragma: no cover
 
     # Make base_vars unmeasurable
-    base_to_unmeasurable_vars = {
-        base_var: assign_custom_measurable_outputs(base_var.owner).outputs[
-            base_var.owner.outputs.index(base_var)
-        ]
-        for base_var in base_vars
-    }
+    base_to_unmeasurable_vars = {base_var: ignore_logprob(base_var) for base_var in base_vars}
 
     def replacement_fn(var, replacements):
         if var in base_to_unmeasurable_vars:
@@ -339,7 +330,7 @@ def find_measurable_dimshuffles(fgraph, node) -> Optional[List[MeasurableDimShuf
         return None  # pragma: no cover
 
     # Make base_vars unmeasurable
-    base_var = assign_custom_measurable_outputs(base_var.owner)
+    base_var = ignore_logprob(base_var)
 
     measurable_dimshuffle = MeasurableDimShuffle(node.op.input_broadcastable, node.op.new_order)(
         base_var

pymc/logprob/transforms.py

@@ -78,11 +78,10 @@
     MeasurableVariable,
     _get_measurable_outputs,
     _logprob,
-    assign_custom_measurable_outputs,
     logprob,
 )
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
-from pymc.logprob.utils import walk_model
+from pymc.logprob.utils import ignore_logprob, walk_model
 
 
 class TransformedVariable(Op):
@@ -549,7 +548,7 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
 
     # Make base_measure outputs unmeasurable
     # This seems to be the only thing preventing nested rewrites from being erased
-    measurable_input = assign_custom_measurable_outputs(measurable_input.owner)
+    measurable_input = ignore_logprob(measurable_input)
 
     scalar_op = node.op.scalar_op
     measurable_input_idx = 0

pymc/logprob/utils.py

@@ -275,6 +275,7 @@ def ignore_logprob(rv: TensorVariable) -> TensorVariable:
     op_type = type(node.op)
     if op_type.__name__.startswith(prefix):
         return rv
+    # By default `assign_custom_measurable_outputs` makes all outputs unmeasurable
     new_node = assign_custom_measurable_outputs(node, type_prefix=prefix)
     return new_node.outputs[node.outputs.index(rv)]