Logprob derivation for switch-encoding graphs

pymc/logprob/censoring.py

@@ -34,22 +34,27 @@
 #   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #   SOFTWARE.
 
-from typing import List, Optional
+from typing import Callable, Container, Generator, Iterable, List, Optional, Set, Tuple
 
 import numpy as np
 import pytensor.tensor as pt
 
-from pytensor.graph.basic import Node
+from pytensor.graph.basic import Node, walk
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import node_rewriter
-from pytensor.scalar.basic import Ceil, Clip, Floor, RoundHalfToEven
+from pytensor.scalar.basic import Ceil, Clip, Floor, RoundHalfToEven, Switch
 from pytensor.scalar.basic import clip as scalar_clip
+from pytensor.scalar.basic import switch as scalar_switch
+from pytensor.tensor.basic import switch as switch
 from pytensor.tensor.math import ceil, clip, floor, round_half_to_even
-from pytensor.tensor.var import TensorConstant
+from pytensor.tensor.random.op import RandomVariable
+from pytensor.tensor.var import TensorConstant, TensorVariable
 
-from pymc.logprob.abstract import MeasurableElemwise, _logcdf, _logprob
+from pymc.logprob.abstract import MeasurableElemwise, _logcdf, _logprob, _logprob_helper
+from pymc.logprob.binary import MeasurableBitwise
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
-from pymc.logprob.utils import CheckParameterValue
+from pymc.logprob.utils import CheckParameterValue, check_potential_measurability
+from pymc.pytensorf import replace_rvs_by_values
 
 
 class MeasurableClip(MeasurableElemwise):
@@ -237,3 +242,171 @@
     from pymc.math import logdiffexp
 
     return logdiffexp(logcdf_upper, logcdf_lower)
+
+
+class MeasurableSwitchEncoding(MeasurableElemwise):
+    """A placeholder used to specify the log-likelihood for a encoded RV sub-graph."""
+
+    valid_scalar_types = (Switch,)
+
+    def __init__(self, measurable_branches):
+        super().__init__(scalar_switch)
+        self.__props__ = super().__props__ + ("measurable_branches",)
+        self.measurable_branches = measurable_branches
+        # number of measurable branches to facilitate correct logprob calculation
+
+
+@node_rewriter(tracks=[switch])
+def find_measurable_switch_encoding(
+    fgraph: FunctionGraph, node: Node
+) -> Optional[List[MeasurableSwitchEncoding]]:
+    rv_map_feature: Optional[PreserveRVMappings] = getattr(fgraph, "preserve_rv_mappings", None)
+
+    if rv_map_feature is None:
+        return None  # pragma: no cover
+
+    valued_rvs = rv_map_feature.rv_values.keys()
+
+    switch_condn, *components = node.inputs
+
+    # broadcasting of switch condition is not supported
+    if switch_condn.type.broadcastable != node.outputs[0].type.broadcastable:
+        return None
+
+    measurable_comp_list = [
+        idx
+        for idx, component in enumerate(components)
+        if check_potential_measurability([component], valued_rvs)
+    ]
+
+    # this automatically checks the measurability of the switch condition and converts switch to MeasurableSwitch
+    if rv_map_feature.request_measurable([switch_condn]) != [switch_condn]:
+        return None
+
+    [base_var] = rv_map_feature.request_measurable([switch_condn.owner.inputs[0]])
+    if base_var.dtype.startswith("int"):
+        return None
+
+    # default number of measurable branches is zero
+    measurable_switch_encoding = MeasurableSwitchEncoding(measurable_branches=0)
+
+    # Maximum one branch allowed to be measurable
+    if len(measurable_comp_list) > 1:
+        return None
+
+    # If at least one of the branches is measurable
+    if len(measurable_comp_list) == 1:
+        measurable_comp_idx = measurable_comp_list[0]
+        measurable_component = components[measurable_comp_idx]
+
+        measurable_switch_encoding.measurable_branches = 1
+
+        # broadcasting of the measurable component is not supported
+        if (
+            (measurable_component.type.broadcastable != node.outputs[0].broadcastable)
+            or (not compare_measurability_source([switch_condn, measurable_component], valued_rvs))
+            or (not rv_map_feature.request_measurable([measurable_component]))
+        ):
+            return None
+
+        if measurable_comp_idx == 0:
+            # changing the first branch of switch to always be the encoding
+            inverted_switch = pt.invert(switch_condn)
+
+            bitwise_op = MeasurableBitwise(inverted_switch.owner.op.scalar_op)
+            measurable_inverted_switch = bitwise_op.make_node(switch_condn).default_output()
+            encoded_rv = measurable_switch_encoding.make_node(
+                measurable_inverted_switch, *components[::-1]
+            ).default_output()
+
+            return [encoded_rv]
+
+    encoded_rv = measurable_switch_encoding.make_node(switch_condn, *components).default_output()
+
+    return [encoded_rv]
+
+
+@_logprob.register(MeasurableSwitchEncoding)
+def switch_encoding_logprob(op, values, *inputs, **kwargs):
+    (value,) = values
+
+    switch_condn, *components = inputs
+
+    if op.measurable_branches == 0:
+        logprob = pt.switch(
+            pt.eq(value, components[0]),
+            _logprob_helper(switch_condn, pt.as_tensor(np.array(True)), **kwargs),
+            pt.switch(
+                pt.eq(value, components[1]),
+                _logprob_helper(switch_condn, pt.as_tensor(np.array(False))),
+                -np.inf,
+            ),
+        )
+    else:
+        base_var = components[1]
+
+        logp_first_branch = _logprob_helper(switch_condn, pt.as_tensor(np.array(True)), **kwargs)
+
+        (switch_condn,) = replace_rvs_by_values([switch_condn], rvs_to_values={base_var: value})
+        logprob = pt.switch(
+            pt.eq(value, components[0]),
+            logp_first_branch,
+            pt.switch(
+                pt.invert(switch_condn),
+                _logprob_helper(base_var, value, **kwargs),
+                -np.inf,
+            ),
+        )
+
+    return logprob
+
+
+measurable_ir_rewrites_db.register(
+    "find_measurable_switch_encoding", find_measurable_switch_encoding, "basic", "censoring"
+)
+
+
+def compare_measurability_source(
+    inputs: Tuple[TensorVariable], valued_rvs: Container[TensorVariable]
+) -> bool:
+    ancestor_var_set = set()
+
+    # retrieve the source of measurability for all elements in 'inputs' separately.
+    for inp in inputs:
+        for ancestor_var in walk_model(
+            [inp],
+            walk_past_rvs=False,
+            stop_at_vars=set(valued_rvs),
+        ):
+            if (
+                ancestor_var.owner
+                and isinstance(ancestor_var.owner.op, RandomVariable)
+                and ancestor_var not in valued_rvs
+            ):
+                ancestor_var_set.add(ancestor_var)
+
+    return len(ancestor_var_set) == 1
+
+
+def walk_model(
+    graphs: Iterable[TensorVariable],
+    walk_past_rvs: bool = False,
+    stop_at_vars: Optional[Set[TensorVariable]] = None,
+    expand_fn: Callable[[TensorVariable], List[TensorVariable]] = lambda var: [],
+) -> Generator[TensorVariable, None, None]:
+    if stop_at_vars is None:
+        stop_at_vars = set()
+
+    def expand(var: TensorVariable, stop_at_vars=stop_at_vars) -> List[TensorVariable]:
+        new_vars = expand_fn(var)
+
+        if (
+            var.owner
+            and (walk_past_rvs or not isinstance(var.owner.op, RandomVariable))
+            and (var not in stop_at_vars)
+        ):
+            new_vars.extend(reversed(var.owner.inputs))
+
+        return new_vars
+
+    yield from walk(graphs, expand, False)

tests/logprob/test_censoring.py

@@ -43,6 +43,9 @@
 
 from pymc import logp
 from pymc.logprob import conditional_logp
+from pymc.logprob.abstract import MeasurableVariable
+from pymc.logprob.censoring import MeasurableSwitchEncoding
+from pymc.logprob.rewriting import construct_ir_fgraph
 from pymc.logprob.transforms import LogTransform, TransformValuesRewrite
 from pymc.testing import assert_no_rvs
 
@@ -261,3 +264,77 @@ def test_rounding(rounding_op):
         logprob.eval({xr_vv: test_value}),
         expected_logp,
     )
+
+
+def test_switch_encoding_both_branches():
+    x_rv = pt.random.normal(0.5, 1)
+    y_rv = pt.switch(x_rv < 0.3, 1.0, 2.0)
+
+    y_vv = y_rv.clone()
+    ref_scipy = st.norm(0.5, 1)
+
+    logprob = logp(y_rv, y_vv)
+    logp_fn = pytensor.function([y_vv], logprob)
+
+    assert logp_fn(1.5) == -np.inf
+
+    assert np.isclose(logp_fn(1), ref_scipy.logcdf(0.3))
+    assert np.isclose(logp_fn(2), ref_scipy.logsf(0.3))
+
+
+@pytest.mark.parametrize(
+    "measurable_idx, test_values, exp_logp",
+    [
+        (1, (0.9, 1, 1.5), (-np.inf, st.norm(0.5, 1).logcdf(1), st.norm(0.5, 1).logpdf(1.5))),
+        (0, (1.5, 1, 0.9), (-np.inf, st.norm(0.5, 1).logsf(1), st.norm(0.5, 1).logpdf(0.9))),
+    ],
+)
+def test_switch_encoding_one_branch_measurable(measurable_idx, test_values, exp_logp):
+    x_rv = pt.random.normal(0.5, 1)
+    branches = (1, x_rv) if measurable_idx == 1 else (x_rv, 1)
+
+    y_rv = pt.switch(x_rv < 1, *branches)
+
+    y_vv = y_rv.clone()
+
+    logprob = logp(y_rv, y_vv)
+
+    logp_fn = pytensor.function([y_vv], logprob)
+
+    for i, j in zip(test_values, exp_logp):
+        assert np.isclose(logp_fn(i), j)
+
+
+def test_switch_encoding_invalid_bcast():
+    x_rv = pt.random.normal(0.5, 1)
+
+    y_rv = pt.switch(x_rv < 0.3, 0.0, 1.0)
+    y_rv_invalid = pt.switch(x_rv < 0.3, [0.0, 0.5], 1.0)
+
+    y_vv = y_rv.clone()
+    y_vv_invalid = y_rv_invalid.clone()
+
+    y_test = 1.0
+
+    assert np.isclose(logp(y_rv, y_vv).eval({y_vv: y_test}), st.norm(0.5, 1).logsf(0.3))
+
+    with pytest.raises(
+        NotImplementedError,
+        match="Logprob method not implemented",
+    ):
+        logp(y_rv_invalid, y_vv_invalid).eval({y_vv_invalid: y_test})
+
+
+def test_switch_encoding_discrete_fail():
+    """We do not support the encoding graphs of discrete RVs yet"""
+    x_rv = pt.random.poisson(2)
+    y_rv = pt.switch(x_rv > 3, x_rv, 1)
+
+    y_vv = x_rv.clone()
+    y_vv_test = 1
+
+    with pytest.raises(
+        NotImplementedError,
+        match="Logprob method not implemented",
+    ):
+        logp(y_rv, y_vv).eval({y_vv: y_vv_test})