Allow logcdf inference of discrete variables

Author: ricardoV94

pymc/logprob/transforms.py

@@ -232,20 +232,20 @@ def measurable_transform_logcdf(op: MeasurableTransform, value, *inputs, **kwarg
     """Compute the log-CDF graph for a `MeasurabeTransform`."""
     other_inputs = list(inputs)
     measurable_input = other_inputs.pop(op.measurable_input_idx)
-
-    # Do not apply rewrite to discrete variables
-    if measurable_input.type.dtype.startswith("int"):
-        raise NotImplementedError("logcdf of transformed discrete variables not implemented")
-
     backward_value = op.transform_elemwise.backward(value, *other_inputs)
 
     # Fail if transformation is not injective
     # A TensorVariable is returned in 1-to-1 inversions, and a tuple in 1-to-many
     if isinstance(backward_value, tuple):
         raise NotImplementedError
 
+    is_discrete = measurable_input.type.dtype.startswith("int")
+
     logcdf = _logcdf_helper(measurable_input, backward_value)
-    logccdf = pt.log1mexp(logcdf)
+    if is_discrete:
+        logccdf = pt.log1mexp(_logcdf_helper(measurable_input, backward_value - 1))
+    else:
+        logccdf = pt.log1mexp(logcdf)
 
     if isinstance(op.scalar_op, MONOTONICALLY_INCREASING_OPS):
         pass
@@ -271,7 +271,6 @@ def measurable_transform_logcdf(op: MeasurableTransform, value, *inputs, **kwarg
 
     # The jacobian is used to ensure a value in the supported domain was provided
     jacobian = op.transform_elemwise.log_jac_det(value, *other_inputs)
-
     return pt.switch(pt.isnan(jacobian), -np.inf, logcdf)
 
 

tests/logprob/test_transforms.py

@@ -693,37 +693,42 @@ def test_not_implemented_discrete_rv_transform():
 
 def test_negated_discrete_rv_transform():
     p = 0.7
-    rv = -Bernoulli.dist(p=p)
+    rv = -Bernoulli.dist(p=p, shape=(4,))
     vv = rv.type()
-    logp_fn = pytensor.function([vv], logp(rv, vv))
 
     # A negated Bernoulli has pmf {p if x == -1; 1-p if x == 0; 0 otherwise}
-    assert logp_fn(-2) == -np.inf
-    np.testing.assert_allclose(logp_fn(-1), np.log(p))
-    np.testing.assert_allclose(logp_fn(0), np.log(1 - p))
-    assert logp_fn(1) == -np.inf
+    logp_fn = pytensor.function([vv], logp(rv, vv))
+    np.testing.assert_allclose(
+        logp_fn([-2, -1, 0, 1]), [-np.inf, np.log(p), np.log(1 - p), -np.inf]
+    )
 
-    # Logcdf and icdf not supported yet
-    for func in (logcdf, icdf):
-        with pytest.raises(NotImplementedError):
-            func(rv, 0)
+    logcdf_fn = pytensor.function([vv], logcdf(rv, vv))
+    np.testing.assert_allclose(logcdf_fn([-2, -1, 0, 1]), [-np.inf, np.log(p), 0, 0])
+
+    with pytest.raises(NotImplementedError):
+        icdf(rv, [-2, -1, 0, 1])
 
 
 def test_shifted_discrete_rv_transform():
     p = 0.7
     rv = Bernoulli.dist(p=p) + 5
     vv = rv.type()
-    rv_logp_fn = pytensor.function([vv], logp(rv, vv))
 
+    rv_logp_fn = pytensor.function([vv], logp(rv, vv))
     assert rv_logp_fn(4) == -np.inf
     np.testing.assert_allclose(rv_logp_fn(5), np.log(1 - p))
     np.testing.assert_allclose(rv_logp_fn(6), np.log(p))
     assert rv_logp_fn(7) == -np.inf
 
-    # Logcdf and icdf not supported yet
-    for func in (logcdf, icdf):
-        with pytest.raises(NotImplementedError):
-            func(rv, 0)
+    rv_logcdf_fn = pytensor.function([vv], logcdf(rv, vv))
+    assert rv_logcdf_fn(4) == -np.inf
+    np.testing.assert_allclose(rv_logcdf_fn(5), np.log(1 - p))
+    np.testing.assert_allclose(rv_logcdf_fn(6), 0)
+    assert rv_logcdf_fn(7) == 0
+
+    # icdf not supported yet
+    with pytest.raises(NotImplementedError):
+        icdf(rv, 0)
 
 
 @pytest.mark.xfail(reason="Check not implemented yet")