Fix logprob of check_and_raise

Author: ricardoV94

.github/workflows/tests.yml

@@ -105,6 +105,7 @@ jobs:
             tests/logprob/test_abstract.py
             tests/logprob/test_basic.py
             tests/logprob/test_binary.py
+            tests/logprob/test_checks.py
             tests/logprob/test_censoring.py
             tests/logprob/test_composite_logprob.py
             tests/logprob/test_cumsum.py

pymc/logprob/checks.py

@@ -39,7 +39,7 @@
 import pytensor.tensor as pt
 
 from pytensor.graph.rewriting.basic import node_rewriter
-from pytensor.raise_op import CheckAndRaise, ExceptionType
+from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor.shape import SpecifyShape
 
 from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
@@ -106,15 +106,18 @@ class MeasurableCheckAndRaise(CheckAndRaise):
 
 
 @_logprob.register(MeasurableCheckAndRaise)
-def logprob_assert(op, values, inner_rv, *assertion, **kwargs):
+def logprob_check_and_raise(op, values, inner_rv, *assertions, **kwargs):
+    from pymc.pytensorf import replace_rvs_by_values
+
     (value,) = values
     # transfer assertion from rv to value
-    value = op(assertion, value)
+    assertions = replace_rvs_by_values(assertions, rvs_to_values={inner_rv: value})
+    value = op(value, *assertions)
     return _logprob_helper(inner_rv, value)
 
 
 @node_rewriter([CheckAndRaise])
-def find_measurable_asserts(fgraph, node) -> Optional[List[MeasurableCheckAndRaise]]:
+def find_measurable_check_and_raise(fgraph, node) -> Optional[List[MeasurableCheckAndRaise]]:
     r"""Finds `AssertOp`\s for which a `logprob` can be computed."""
 
     if isinstance(node.op, MeasurableCheckAndRaise):
@@ -126,24 +129,19 @@ def find_measurable_asserts(fgraph, node) -> Optional[List[MeasurableCheckAndRai
         return None  # pragma: no cover
 
     base_rv, *conds = node.inputs
+    if not rv_map_feature.request_measurable([base_rv]):
+        return None
 
-    if not (
-        base_rv.owner
-        and isinstance(base_rv.owner.op, MeasurableVariable)
-        and base_rv not in rv_map_feature.rv_values
-    ):
-        return None  # pragma: no cover
-
-    exception_type = ExceptionType()
-    new_op = MeasurableCheckAndRaise(exc_type=exception_type)
+    op = node.op
+    new_op = MeasurableCheckAndRaise(exc_type=op.exc_type, msg=op.msg)
     new_rv = new_op.make_node(base_rv, *conds).default_output()
 
     return [new_rv]
 
 
 measurable_ir_rewrites_db.register(
-    "find_measurable_asserts",
-    find_measurable_asserts,
+    "find_measurable_check_and_raise",
+    find_measurable_check_and_raise,
     "basic",
     "assert",
 )

tests/logprob/test_checks.py

@@ -81,7 +81,7 @@ def test_assert_logprob():
     rv = pt.random.normal()
     assert_op = Assert("Test assert")
     # Example: Add assert that rv must be positive
-    assert_rv = assert_op(rv > 0, rv)
+    assert_rv = assert_op(rv, rv > 0)
     assert_rv.name = "assert_rv"
 
     assert_vv = assert_rv.clone()
@@ -90,8 +90,10 @@ def test_assert_logprob():
     # Check valid value is correct and doesn't raise
     # Since here the value to the rv satisfies the condition, no error is raised.
     valid_value = 3.0
-    with pytest.raises(AssertionError, match="Test assert"):
-        assert_logp.eval({assert_vv: valid_value})
+    np.testing.assert_allclose(
+        assert_logp.eval({assert_vv: valid_value}),
+        stats.norm.logpdf(valid_value),
+    )
 
     # Check invalid value
     # Since here the value to the rv is negative, an exception is raised as the condition is not met