Allow logcdf and icdf inference

Author: ricardoV94

pymc/logprob/basic.py

@@ -80,13 +80,25 @@ def logp(rv: TensorVariable, value: TensorLike, **kwargs) -> TensorVariable:
 def logcdf(rv: TensorVariable, value: TensorLike, **kwargs) -> TensorVariable:
     """Create a graph for the log-CDF of a Random Variable."""
     value = pt.as_tensor_variable(value, dtype=rv.dtype)
-    return _logcdf_helper(rv, value, **kwargs)
+    try:
+        return _logcdf_helper(rv, value, **kwargs)
+    except NotImplementedError:
+        # Try to rewrite rv
+        fgraph, rv_values, _ = construct_ir_fgraph({rv: value})
+        [ir_rv] = fgraph.outputs
+        return _logcdf_helper(ir_rv, value, **kwargs)
 
 
 def icdf(rv: TensorVariable, value: TensorLike, **kwargs) -> TensorVariable:
     """Create a graph for the inverse CDF of a  Random Variable."""
-    value = pt.as_tensor_variable(value)
-    return _icdf_helper(rv, value, **kwargs)
+    value = pt.as_tensor_variable(value, dtype=rv.dtype)
+    try:
+        return _icdf_helper(rv, value, **kwargs)
+    except NotImplementedError:
+        # Try to rewrite rv
+        fgraph, rv_values, _ = construct_ir_fgraph({rv: value})
+        [ir_rv] = fgraph.outputs
+        return _icdf_helper(ir_rv, value, **kwargs)
 
 
 def factorized_joint_logprob(

pymc/logprob/transforms.py

@@ -77,6 +77,10 @@
     MeasurableElemwise,
     MeasurableVariable,
     _get_measurable_outputs,
+    _icdf,
+    _icdf_helper,
+    _logcdf,
+    _logcdf_helper,
     _logprob,
     _logprob_helper,
 )
@@ -390,6 +394,38 @@ def measurable_transform_logprob(op: MeasurableTransform, values, *inputs, **kwa
     return pt.switch(pt.isnan(jacobian), -np.inf, input_logprob + jacobian)
 
 
+@_logcdf.register(MeasurableTransform)
+def measurable_transform_logcdf(op: MeasurableTransform, value, *inputs, **kwargs):
+    """Compute the log-CDF graph for a `MeasurabeTransform`."""
+    other_inputs = list(inputs)
+    measurable_input = other_inputs.pop(op.measurable_input_idx)
+
+    backward_value = op.transform_elemwise.backward(value, *other_inputs)
+
+    # Some transformations, like squaring may produce multiple backward values
+    if isinstance(backward_value, tuple):
+        raise NotImplementedError
+
+    input_logcdf = _logcdf_helper(measurable_input, backward_value)
+
+    # 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, input_logcdf)
+
+
+@_icdf.register(MeasurableTransform)
+def measurable_transform_icdf(op: MeasurableTransform, value, *inputs, **kwargs):
+    """Compute the inverse CDF graph for a `MeasurabeTransform`."""
+    other_inputs = list(inputs)
+    measurable_input = other_inputs.pop(op.measurable_input_idx)
+
+    input_icdf = _icdf_helper(measurable_input, value)
+    icdf = op.transform_elemwise.forward(input_icdf, *other_inputs)
+
+    return icdf
+
+
 @node_rewriter([reciprocal])
 def measurable_reciprocal_to_power(fgraph, node):
     """Convert reciprocal of `MeasurableVariable`s to power."""

tests/logprob/test_basic.py

@@ -432,8 +432,8 @@ def test_probability_direct_dispatch(func, scipy_func):
     "func, scipy_func, test_value",
     [
         (logp, "logpdf", 5.0),
-        pytest.param(logcdf, "logcdf", 5.0, marks=pytest.mark.xfail(raises=NotImplementedError)),
-        pytest.param(icdf, "ppf", 0.7, marks=pytest.mark.xfail(raises=NotImplementedError)),
+        (logcdf, "logcdf", 5.0),
+        (icdf, "ppf", 0.7),
     ],
 )
 def test_probability_inference(func, scipy_func, test_value):