Allow composition of interdependent container variables

Join/MakeVector/IfElse can output multiple interdependent variables. These are potentially measurable because in the logp each output is given a distinct value variable. However, this isn't known during the IR rewrites.

To circumvent this issue, we run an inner IR rewrite after giving dummy value variables to each output

Author: ricardoV94

pymc/logprob/binary.py

@@ -60,7 +60,7 @@ def find_measurable_comparisons(
     const = node.inputs[(measurable_var_idx + 1) % 2]
 
     # check for potential measurability of const
-    if not check_potential_measurability([const], rv_map_feature):
+    if check_potential_measurability([const], rv_map_feature.rv_values.keys()):
         return None
 
     node_scalar_op = node.op.scalar_op

pymc/logprob/mixture.py

@@ -66,11 +66,13 @@
 from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
 from pymc.logprob.rewriting import (
     PreserveRVMappings,
+    assume_measured_ir_outputs,
     local_lift_DiracDelta,
     measurable_ir_rewrites_db,
     subtensor_ops,
 )
 from pymc.logprob.tensor import naive_bcast_rv_lift
+from pymc.logprob.utils import check_potential_measurability
 
 
 def is_newaxis(x):
@@ -453,19 +455,66 @@ class MeasurableIfElse(IfElse):
 MeasurableVariable.register(MeasurableIfElse)
 
 
+@node_rewriter([IfElse])
+def useless_ifelse_outputs(fgraph, node):
+    """Remove outputs that are shared across the IfElse branches."""
+    # TODO: This should be a PyTensor canonicalization
+    op = node.op
+    if_var, *inputs = node.inputs
+    shared_inputs = set(inputs[op.n_outs :]).intersection(inputs[: op.n_outs])
+    if not shared_inputs:
+        return None
+
+    replacements = {}
+    for shared_inp in shared_inputs:
+        idx = inputs.index(shared_inp)
+        replacements[node.outputs[idx]] = shared_inp
+
+    # IfElse isn't needed at all
+    if len(shared_inputs) == op.n_outs:
+        return replacements
+
+    # Create subset IfElse with remaining nodes
+    remaining_inputs = [inp for inp in inputs if inp not in shared_inputs]
+    new_outs = (
+        IfElse(n_outs=len(remaining_inputs) // 2).make_node(if_var, *remaining_inputs).outputs
+    )
+    for inp, new_out in zip(remaining_inputs, new_outs):
+        idx = inputs.index(inp)
+        replacements[node.outputs[idx]] = new_out
+
+    return replacements
+
+
 @node_rewriter([IfElse])
 def find_measurable_ifelse_mixture(fgraph, node):
     rv_map_feature: Optional[PreserveRVMappings] = getattr(fgraph, "preserve_rv_mappings", None)
 
     if rv_map_feature is None:
         return None  # pragma: no cover
 
+    op = node.op
     if_var, *base_rvs = node.inputs
 
-    if rv_map_feature.request_measurable(base_rvs) != base_rvs:
+    valued_rvs = rv_map_feature.rv_values.keys()
+    if not all(check_potential_measurability([base_var], valued_rvs) for base_var in base_rvs):
         return None
 
-    return MeasurableIfElse(n_outs=node.op.n_outs).make_node(if_var, *base_rvs).outputs
+    base_rvs = assume_measured_ir_outputs(valued_rvs, base_rvs)
+    if len(base_rvs) != op.n_outs * 2:
+        return None
+    if not all(var.owner and isinstance(var.owner.op, MeasurableVariable) for var in base_rvs):
+        return None
+
+    return MeasurableIfElse(n_outs=op.n_outs).make_node(if_var, *base_rvs).outputs
+
+
+measurable_ir_rewrites_db.register(
+    "useless_ifelse_outputs",
+    useless_ifelse_outputs,
+    "basic",
+    "mixture",
+)
 
 
 measurable_ir_rewrites_db.register(

pymc/logprob/rewriting.py

@@ -42,9 +42,16 @@
 
 from pytensor import config
 from pytensor.compile.mode import optdb
-from pytensor.graph.basic import Constant, Variable, ancestors, io_toposort
+from pytensor.graph.basic import (
+    Constant,
+    Variable,
+    ancestors,
+    io_toposort,
+    truncated_graph_inputs,
+)
 from pytensor.graph.features import Feature
 from pytensor.graph.fg import FunctionGraph
+from pytensor.graph.replace import clone_replace
 from pytensor.graph.rewriting.basic import (
     ChangeTracker,
     EquilibriumGraphRewriter,
@@ -461,3 +468,34 @@ 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)
+
+
+def assume_measured_ir_outputs(
+    inputs: Sequence[TensorVariable], outputs: Sequence[TensorVariable]
+) -> Sequence[TensorVariable]:
+    """Run IR rewrite assuming each output is measured.
+
+    IR variables could depend on each other in a way that looks unmeasurable without a value variable assigned to each.
+    For instance `join([add(x, z), z])` is a potentially measurable join, but `add(x, z)` can look unmeasurable
+    because neither `x` and `z` are valued in the IR representation.
+    This helper runs an inner ir rewrite after giving each output a dummy value variable.
+    We replace inputs by dummies and then undo it so that any dependency on outer variables is preserved.
+    """
+    # Replace inputs by dummy variables
+    replaced_inputs = {
+        var: var.type()
+        for var in truncated_graph_inputs(outputs, ancestors_to_include=inputs)
+        if var in inputs
+    }
+    cloned_outputs = clone_replace(outputs, replace=replaced_inputs)
+
+    dummy_rv_values = {base_var: base_var.type() for base_var in cloned_outputs}
+    fgraph, *_ = construct_ir_fgraph(dummy_rv_values)
+
+    # Replace dummy variables by inputs
+    fgraph.replace_all(
+        tuple((repl, orig) for orig, repl in replaced_inputs.items()),
+        import_missing=True,
+    )
+
+    return fgraph.outputs

pymc/logprob/tensor.py

@@ -51,7 +51,12 @@
 )
 
 from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
-from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
+from pymc.logprob.rewriting import (
+    PreserveRVMappings,
+    assume_measured_ir_outputs,
+    measurable_ir_rewrites_db,
+)
+from pymc.logprob.utils import check_potential_measurability
 
 
 @node_rewriter([BroadcastTo])
@@ -213,7 +218,12 @@ def find_measurable_stacks(
     else:
         base_vars = node.inputs
 
-    if rv_map_feature.request_measurable(base_vars) != base_vars:
+    valued_rvs = rv_map_feature.rv_values.keys()
+    if not all(check_potential_measurability([base_var], valued_rvs) for base_var in base_vars):
+        return None
+
+    base_vars = assume_measured_ir_outputs(valued_rvs, base_vars)
+    if not all(var.owner and isinstance(var.owner.op, MeasurableVariable) for var in base_vars):
         return None
 
     if is_join:

pymc/logprob/transforms.py

@@ -529,7 +529,7 @@ def find_measurable_transforms(fgraph: FunctionGraph, node: Node) -> Optional[Li
     # would be invalid
     other_inputs = tuple(inp for inp in node.inputs if inp is not measurable_input)
 
-    if not check_potential_measurability(other_inputs, rv_map_feature):
+    if check_potential_measurability(other_inputs, rv_map_feature.rv_values.keys()):
         return None
 
     scalar_op = node.op.scalar_op

pymc/logprob/utils.py

@@ -38,6 +38,7 @@
 
 from typing import (
     Callable,
+    Container,
     Dict,
     Generator,
     Iterable,
@@ -210,22 +211,24 @@ def indices_from_subtensor(idx_list, indices):
     )
 
 
-def check_potential_measurability(inputs: Tuple[TensorVariable], rv_map_feature):
+def check_potential_measurability(
+    inputs: Tuple[TensorVariable], valued_rvs: Container[TensorVariable]
+) -> bool:
     if any(
-        ancestor_node
-        for ancestor_node in walk_model(
+        ancestor_var
+        for ancestor_var in walk_model(
             inputs,
             walk_past_rvs=False,
-            stop_at_vars=set(rv_map_feature.rv_values),
+            stop_at_vars=set(valued_rvs),
         )
         if (
-            ancestor_node.owner
-            and isinstance(ancestor_node.owner.op, MeasurableVariable)
-            and ancestor_node not in rv_map_feature.rv_values
+            ancestor_var.owner
+            and isinstance(ancestor_var.owner.op, MeasurableVariable)
+            and ancestor_var not in valued_rvs
         )
     ):
-        return None
-    return True
+        return True
+    return False
 
 
 class ParameterValueError(ValueError):

tests/logprob/test_composite_logprob.py

@@ -40,7 +40,7 @@
 import pytest
 import scipy.stats as st
 
-from pymc import logp
+from pymc import draw, logp
 from pymc.logprob.abstract import MeasurableVariable
 from pymc.logprob.basic import factorized_joint_logprob
 from pymc.logprob.censoring import MeasurableClip
@@ -218,3 +218,51 @@ def test_affine_log_transform_rv():
         logp_fn(a_val, b_val, y_val),
         st.norm(a_val, b_val).logpdf(y_val),
     )
+
+
+@pytest.mark.parametrize("reverse", (False, True))
+def test_affine_join_interdependent(reverse):
+    x = pt.random.normal(name="x")
+    y_rvs = []
+    prev_rv = x
+    for i in range(3):
+        next_rv = pt.exp(prev_rv + pt.random.beta(3, 1, name=f"y{i}", size=(1, 2)))
+        y_rvs.append(next_rv)
+        prev_rv = next_rv
+
+    if reverse:
+        y_rvs = y_rvs[::-1]
+
+    ys = pt.concatenate(y_rvs, axis=0)
+    ys.name = "ys"
+
+    x_vv = x.clone()
+    ys_vv = ys.clone()
+
+    logp = factorized_joint_logprob({x: x_vv, ys: ys_vv})
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
+    assert_no_rvs(logp_combined)
+
+    y0_vv = y_rvs[0].clone()
+    y1_vv = y_rvs[1].clone()
+    y2_vv = y_rvs[2].clone()
+
+    ref_logp = factorized_joint_logprob(
+        {x: x_vv, y_rvs[0]: y0_vv, y_rvs[1]: y1_vv, y_rvs[2]: y2_vv}
+    )
+    ref_logp_combined = pt.sum([pt.sum(factor) for factor in ref_logp.values()])
+
+    rng = np.random.default_rng()
+    x_vv_test, ys_vv_test = draw([x, ys], random_seed=rng)
+    ys_vv_test = rng.normal(size=(3, 2))
+    np.testing.assert_allclose(
+        logp_combined.eval({x_vv: x_vv_test, ys_vv: ys_vv_test}),
+        ref_logp_combined.eval(
+            {
+                x_vv: x_vv_test,
+                y0_vv: ys_vv_test[0:1],
+                y1_vv: ys_vv_test[1:2],
+                y2_vv: ys_vv_test[2:3],
+            }
+        ),
+    )

tests/logprob/test_mixture.py

@@ -984,7 +984,7 @@ def test_ifelse_mixture_multiple_components():
 
     if_var = pt.scalar("if_var", dtype="bool")
     comp_then1 = pt.random.normal(size=(2,), name="comp_true1")
-    comp_then2 = pt.random.normal(comp_then1, size=(2, 2), name="comp_then2")
+    comp_then2 = comp_then1 + pt.random.normal(size=(2, 2), name="comp_then2")
     comp_else1 = pt.random.halfnormal(size=(4,), name="comp_else1")
     comp_else2 = pt.random.halfnormal(size=(4, 4), name="comp_else2")
 

tests/logprob/test_utils.py

@@ -51,6 +51,7 @@
 from pymc.logprob.basic import joint_logp, logp
 from pymc.logprob.utils import (
     ParameterValueError,
+    check_potential_measurability,
     dirac_delta,
     rvs_to_value_vars,
     walk_model,
@@ -194,3 +195,17 @@ def scipy_logprob(obs, c):
         return 0.0 if obs == c else -np.inf
 
     scipy_logprob_tester(x, obs, dist_params, test_fn=scipy_logprob)
+
+
+def test_check_potential_measurability():
+    x1 = pt.random.normal()
+    x2 = pt.random.normal()
+    x3 = pt.scalar("x3")
+    y = pt.exp(x1 + x2 + x3)
+
+    # In the first three cases, y is potentially measurable, because it has at least on unvalued RV input
+    assert check_potential_measurability([y], {})
+    assert check_potential_measurability([y], {x1})
+    assert check_potential_measurability([y], {x2})
+    # y is not potentially measurable because both RV inputs are valued
+    assert not check_potential_measurability([y], {x1, x2})