Logprob derivation for switch-encoding graphs

pymc/logprob/censoring.py

@@ -54,6 +54,7 @@
 from pymc.logprob.binary import MeasurableBitwise
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
 from pymc.logprob.utils import CheckParameterValue, check_potential_measurability
+from pymc.pytensorf import replace_rvs_by_values
 
 
 class MeasurableClip(MeasurableElemwise):
@@ -247,6 +248,8 @@
     """A placeholder used to specify the log-likelihood for a encoded RV sub-graph."""
 
     valid_scalar_types = (Switch,)
+    # number of measurable branches to facilitate correct logprob calculation
+    measurable_branches = 0
 
 
 measurable_switch_encoding = MeasurableSwitchEncoding(scalar_switch)
@@ -256,88 +259,102 @@
 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
 
     # 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
 
-    # Right now, this only works for switch with both encoding branches.
-    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,
-        ),
-    )
-
-    # TODO: Calculate logprob for switch with one measurable component If RV is discrete,
-    #  give preference over encoding.
+    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]  # there needs to be a better way to obtain the base variable.
+
+        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(
@@ -348,23 +365,23 @@
 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(
@@ -373,19 +390,19 @@
     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

@@ -279,18 +279,17 @@ def test_switch_encoding_both_branches():
     assert np.isclose(logp_fn(2), ref_scipy.logsf(0.3))
 
 
-@pytest.mark.skip(reason="Logprob calculation for measurable branches not added")
 def test_switch_encoding_second_branch_measurable():
     x_rv = pt.random.normal(0.5, 1)
-    y_rv = pt.switch(x_rv < 0.3, 1, x_rv)
+    y_rv = pt.switch(x_rv < 1, 1, x_rv)
 
     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(3) == -np.inf
+    assert logp_fn(0.5) == -np.inf
 
-    assert np.isclose(logp_fn(1), ref_scipy.logcdf(0.3))
-    assert np.isclose(logp_fn(0.2), -np.inf)
+    assert np.isclose(logp_fn(1), ref_scipy.logcdf(1))
+    assert np.isclose(logp_fn(1.2), ref_scipy.logpdf(1.2))