Explicitly remove AePPL IR from conditional_logprob results

Author: brandonwillard

aeppl/abstract.py

@@ -144,6 +144,7 @@ class ValuedVariable(Op):
     directly.  An example is `BroadcastTo` lifting through `RandomVariable`\s.
     """
 
+    __props__ = ()
     default_output = 0
     view_map = {0: [0]}
 

aeppl/joint_logprob.py

@@ -8,7 +8,7 @@
 
 from aeppl.abstract import ValuedVariable, get_measurable_outputs
 from aeppl.logprob import _logprob
-from aeppl.rewriting import construct_ir_fgraph
+from aeppl.rewriting import construct_ir_fgraph, ir_cleanup_db
 
 if TYPE_CHECKING:
     from aesara.graph.basic import Apply, Variable
@@ -219,11 +219,9 @@ def conditional_logprob(
         #     for node in io_toposort(graph_inputs([rv_logprobs]), outputs):
         #         compute_test_value(node)
 
-    # Replace `ValuedVariable`s with their values
+    # Remove unneeded IR elements from the graph
     rv_logprobs_fg = FunctionGraph(outputs=tuple(logprob_vars.values()), clone=False)
-    rv_logprobs_fg.replace_all(
-        tuple((valued_var, valued_var.owner.inputs[1]) for valued_var in fgraph.outputs)
-    )
+    ir_cleanup_db.query("+basic").rewrite(rv_logprobs_fg)
 
     return logprob_vars, value_vars
 

aeppl/rewriting.py

@@ -5,7 +5,7 @@
 from aesara.graph.basic import Apply, Variable
 from aesara.graph.features import Feature
 from aesara.graph.fg import FunctionGraph
-from aesara.graph.rewriting.basic import GraphRewriter, node_rewriter
+from aesara.graph.rewriting.basic import GraphRewriter, in2out, node_rewriter
 from aesara.graph.rewriting.db import EquilibriumDB, RewriteDatabaseQuery, SequenceDB
 from aesara.tensor.elemwise import DimShuffle, Elemwise
 from aesara.tensor.extra_ops import BroadcastTo
@@ -277,3 +277,18 @@ def construct_ir_fgraph(
     fgraph.replace_all(new_to_old, reason="undo-unvalued-measurables")
 
     return fgraph, rv_value_clones, memo
+
+
+@register_useless
+@node_rewriter([ValuedVariable])
+def remove_ValuedVariable(fgraph, node):
+    return [node.inputs[1]]
+
+
+ir_cleanup_db = SequenceDB()
+ir_cleanup_db.name = "ir_cleanup_db"
+ir_cleanup_db.register(
+    "remove-intermediate-ir",
+    in2out(local_remove_DiracDelta, remove_ValuedVariable),
+    "basic",
+)

aeppl/transforms.py

@@ -24,7 +24,7 @@
     valued_variable,
 )
 from aeppl.logprob import _logprob, logprob
-from aeppl.rewriting import measurable_ir_rewrites_db
+from aeppl.rewriting import ir_cleanup_db, measurable_ir_rewrites_db
 
 if TYPE_CHECKING:
     from aesara.graph.rewriting.basic import NodeRewriter
@@ -75,6 +75,7 @@ class TransformedVariable(Op):
 
     """
 
+    __props__ = ()
     view_map = {0: [0]}
 
     def make_node(self, tran_value: TensorVariable, value: TensorVariable):
@@ -101,8 +102,12 @@ def grad(self, args, g_outs):
 @register_useless
 @node_rewriter([TransformedVariable])
 def remove_TransformedVariables(fgraph, node):
-    if isinstance(node.op, TransformedVariable):
-        return [node.inputs[0]]
+    return [node.inputs[0]]
+
+
+ir_cleanup_db.register(
+    "remove-TransformedVariables", in2out(remove_TransformedVariables), "basic"
+)
 
 
 class RVTransform(abc.ABC):

tests/test_abstract.py

@@ -2,6 +2,7 @@
 import aesara.tensor as at
 import numpy as np
 import pytest
+from aesara.compile.mode import get_default_mode
 from aesara.gradient import NullTypeGradError, grad
 from aesara.tensor.random.basic import NormalRV
 
@@ -113,7 +114,9 @@ def test_valued_variable():
     obs_var = valued_variable(rv_var, rv_vv)
 
     rv_val = np.zeros(3)
-    res = obs_var.eval({rv_var: rv_val, rv_vv: np.ones(3)})
+    mode = get_default_mode().excluding("remove_ValuedVariable")
+    obs_var_fn = aesara.function([rv_var, rv_vv], obs_var, mode=mode)
+    res = obs_var_fn(rv_val, np.ones(3))
     assert np.array_equal(res, rv_val)
 
     with pytest.raises(NullTypeGradError):

tests/test_joint_logprob.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 import scipy.stats.distributions as sp
-from aesara.graph.basic import ancestors, equal_computations
+from aesara.graph.basic import ancestors, applys_between, equal_computations
 from aesara.tensor.subtensor import (
     AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
@@ -13,6 +13,7 @@
     Subtensor,
 )
 
+from aeppl.abstract import ValuedVariable
 from aeppl.joint_logprob import conditional_logprob, joint_logprob
 from aeppl.logprob import logprob
 from aeppl.utils import rvs_to_value_vars, walk_model
@@ -279,3 +280,22 @@ def test_deprecations():
 
     with pytest.warns(DeprecationWarning):
         conditional_logprob(realized={X: x}, warn_missing_rvs=True)
+
+
+def test_no_output_ValuedVariables():
+    srng = at.random.RandomStream(0)
+
+    X_at = at.matrix("X")
+    tau_rv = srng.halfcauchy(1)
+    beta_rv = srng.normal(0, tau_rv, size=X_at.shape[-1])
+
+    eta = X_at @ beta_rv
+    p = at.sigmoid(-eta)
+    Y_rv = srng.bernoulli(p)
+
+    logdensity, vvs = joint_logprob(Y_rv, beta_rv, tau_rv)
+
+    assert not any(
+        isinstance(node.op, ValuedVariable)
+        for node in applys_between(ins=vvs, outs=(logdensity,))
+    )