Remove IR Ops from final logprob graph

Author: ricardoV94

pymc/logprob/basic.py

@@ -64,7 +64,7 @@
     _logprob_helper,
     get_measurable_outputs,
 )
-from pymc.logprob.rewriting import construct_ir_fgraph
+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
 
@@ -107,6 +107,7 @@ def logp(
         fgraph, _, _ = construct_ir_fgraph({rv: value})
         [(ir_rv, ir_value)] = fgraph.preserve_rv_mappings.rv_values.items()
         expr = _logprob_helper(ir_rv, ir_value, **kwargs)
+        cleanup_ir([expr])
         if warn_missing_rvs:
             _warn_rvs_in_inferred_graph(expr)
         return expr
@@ -124,6 +125,7 @@ def logcdf(
         fgraph, rv_values, _ = construct_ir_fgraph({rv: value})
         [ir_rv] = fgraph.outputs
         expr = _logcdf_helper(ir_rv, value, **kwargs)
+        cleanup_ir([expr])
         if warn_missing_rvs:
             _warn_rvs_in_inferred_graph(expr)
         return expr
@@ -141,6 +143,7 @@ def icdf(
         fgraph, rv_values, _ = construct_ir_fgraph({rv: value})
         [ir_rv] = fgraph.outputs
         expr = _icdf_helper(ir_rv, value, **kwargs)
+        cleanup_ir([expr])
         if warn_missing_rvs:
             _warn_rvs_in_inferred_graph(expr)
         return expr
@@ -321,6 +324,8 @@ def factorized_joint_logprob(
             f"The logprob terms of the following value variables could not be derived: {missing_value_terms}"
         )
 
+    cleanup_ir(logprob_vars.values())
+
     return logprob_vars
 
 

pymc/logprob/rewriting.py

@@ -34,7 +34,7 @@
 #   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #   SOFTWARE.
 
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Sequence, Tuple
 
 import pytensor.tensor as pt
 
@@ -43,11 +43,17 @@
 from pytensor.graph.features import Feature
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import GraphRewriter, node_rewriter
-from pytensor.graph.rewriting.db import EquilibriumDB, RewriteDatabaseQuery, SequenceDB
+from pytensor.graph.rewriting.db import (
+    EquilibriumDB,
+    LocalGroupDB,
+    RewriteDatabaseQuery,
+    SequenceDB,
+    TopoDB,
+)
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.extra_ops import BroadcastTo
 from pytensor.tensor.random.rewriting import local_subtensor_rv_lift
-from pytensor.tensor.rewriting.basic import register_canonicalize, register_useless
+from pytensor.tensor.rewriting.basic import register_canonicalize
 from pytensor.tensor.rewriting.shape import ShapeFeature
 from pytensor.tensor.subtensor import (
     AdvancedIncSubtensor,
@@ -191,9 +197,8 @@ def local_lift_DiracDelta(fgraph, node):
     return new_node.outputs
 
 
-@register_useless
-@node_rewriter((DiracDelta,))
-def local_remove_DiracDelta(fgraph, node):
+@node_rewriter([DiracDelta])
+def remove_DiracDelta(fgraph, node):
     r"""Remove `DiracDelta`\s."""
     dd_val = node.inputs[0]
     return [dd_val]
@@ -270,6 +275,17 @@ def incsubtensor_rv_replace(fgraph, node):
 
 logprob_rewrites_db.register("post-canonicalize", optdb.query("+canonicalize"), "basic")
 
+# Rewrites that remove IR Ops
+cleanup_ir_rewrites_db = LocalGroupDB()
+cleanup_ir_rewrites_db.name = "cleanup_ir_rewrites_db"
+logprob_rewrites_db.register(
+    "cleanup_ir_rewrites",
+    TopoDB(cleanup_ir_rewrites_db, order="out_to_in", ignore_newtrees=True, failure_callback=None),
+    "cleanup",
+)
+
+cleanup_ir_rewrites_db.register("remove_DiracDelta", remove_DiracDelta, "cleanup")
+
 
 def construct_ir_fgraph(
     rv_values: Dict[Variable, Variable],
@@ -351,3 +367,9 @@ def construct_ir_fgraph(
         fgraph.replace_all(new_to_old)
 
     return fgraph, rv_values, memo
+
+
+def cleanup_ir(vars: Sequence[Variable]) -> None:
+    fgraph = FunctionGraph(outputs=vars, clone=False)
+    ir_rewriter = logprob_rewrites_db.query(RewriteDatabaseQuery(include=["cleanup"]))
+    ir_rewriter.rewrite(fgraph)

pymc/logprob/transforms.py

@@ -88,11 +88,6 @@
     tanh,
     true_div,
 )
-from pytensor.tensor.rewriting.basic import (
-    register_specialize,
-    register_stabilize,
-    register_useless,
-)
 from pytensor.tensor.var import TensorVariable
 
 from pymc.logprob.abstract import (
@@ -106,7 +101,11 @@
     _logprob,
     _logprob_helper,
 )
-from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
+from pymc.logprob.rewriting import (
+    PreserveRVMappings,
+    cleanup_ir_rewrites_db,
+    measurable_ir_rewrites_db,
+)
 from pymc.logprob.utils import check_potential_measurability, ignore_logprob
 
 
@@ -134,15 +133,20 @@ def grad(self, args, g_outs):
 transformed_variable = TransformedVariable()
 
 
-@register_specialize
-@register_stabilize
-@register_useless
 @node_rewriter([TransformedVariable])
 def remove_TransformedVariables(fgraph, node):
     if isinstance(node.op, TransformedVariable):
         return [node.inputs[0]]
 
 
+cleanup_ir_rewrites_db.register(
+    "remove_TransformedVariables",
+    remove_TransformedVariables,
+    "cleanup",
+    "transform",
+)
+
+
 class RVTransform(abc.ABC):
     ndim_supp = None
 

tests/logprob/test_rewriting.py

@@ -40,6 +40,7 @@
 import pytest
 import scipy.stats.distributions as sp
 
+from pytensor.graph import ancestors
 from pytensor.graph.rewriting.basic import in2out
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
@@ -50,8 +51,10 @@
     Subtensor,
 )
 
+from pymc.distributions.transforms import logodds
 from pymc.logprob.basic import factorized_joint_logprob
-from pymc.logprob.rewriting import local_lift_DiracDelta
+from pymc.logprob.rewriting import cleanup_ir, local_lift_DiracDelta
+from pymc.logprob.transforms import TransformedVariable, TransformValuesRewrite
 from pymc.logprob.utils import DiracDelta, dirac_delta
 
 
@@ -88,10 +91,23 @@ def test_local_lift_DiracDelta():
 
 def test_local_remove_DiracDelta():
     c_at = pt.vector()
-    dd_at = dirac_delta(c_at)
+    dd_at = dirac_delta(c_at) + dirac_delta(5)
+    assert sum(isinstance(v.owner.op, DiracDelta) for v in ancestors([dd_at]) if v.owner) == 2
+
+    cleanup_ir([dd_at])
+    assert not any(isinstance(v.owner.op, DiracDelta) for v in ancestors([dd_at]) if v.owner)
+
+
+def test_local_remove_TransformedVariable():
+    p_rv = pt.random.beta(1, 1, name="p")
+    p_vv = p_rv.clone()
+
+    tr = TransformValuesRewrite({p_vv: logodds})
+    [p_logp] = factorized_joint_logprob({p_rv: p_vv}, extra_rewrites=tr).values()
 
-    fn = pytensor.function([c_at], dd_at)
-    assert not any(isinstance(node.op, DiracDelta) for node in fn.maker.fgraph.toposort())
+    assert not any(
+        isinstance(v.owner.op, TransformedVariable) for v in ancestors([p_logp]) if v.owner
+    )
 
 
 @pytest.mark.parametrize(

tests/logprob/test_transforms.py

@@ -561,6 +561,7 @@ def test_original_values_output_dict():
     assert p_vv in logp_dict
 
 
+@pytest.mark.filterwarnings("error")
 def test_mixture_transform():
     """Make sure that non-`RandomVariable` `MeasurableVariable`s can be transformed.
 
@@ -588,21 +589,17 @@ def test_mixture_transform():
 
     transform_rewrite = TransformValuesRewrite({y_vv: LogTransform()})
 
-    with pytest.warns(None) as record:
-        # This shouldn't raise any warnings
-        logp_trans = factorized_joint_logprob(
-            {Y_rv: y_vv, I_rv: i_vv},
-            extra_rewrites=transform_rewrite,
-            use_jacobian=False,
-        )
-        logp_trans_combined = pt.sum([pt.sum(factor) for factor in logp_trans.values()])
-
-    assert not record.list
+    logp_trans = factorized_joint_logprob(
+        {Y_rv: y_vv, I_rv: i_vv},
+        extra_rewrites=transform_rewrite,
+        use_jacobian=False,
+    )
+    logp_trans_combined = pt.sum([pt.sum(factor) for factor in logp_trans.values()])
 
     # The untransformed graph should be the same as the transformed graph after
     # replacing the `Y_rv` value variable with a transformed version of itself
     logp_nt_fg = FunctionGraph(outputs=[logp_no_trans_comb], clone=False)
-    y_trans = transformed_variable(pt.exp(y_vv), y_vv)
+    y_trans = pt.exp(y_vv)
     y_trans.name = "y_log"
     logp_nt_fg.replace(y_vv, y_trans)
     logp_nt = logp_nt_fg.outputs[0]

tests/test_model.py

@@ -1625,3 +1625,12 @@ def test_invalid_observed_value(self, capfd):
             "Some of the observed values of variable y are associated with a non-finite logp" in out
         )
         assert "value = 0.53 -> logp = -inf" in out
+
+
+def test_model_logp_fast_compile():
+    # Issue #5618
+    with pm.Model() as m:
+        pm.Dirichlet("a", np.ones(3))
+
+    with pytensor.config.change_flags(mode="FAST_COMPILE"):
+        assert m.point_logps() == {"a": -1.5}