Move try/except fallback into _logccdf_helper

Centralizes the fallback logic so callers don't need to handle it.
The helper now tries the stable _logccdf first and automatically
falls back to log1mexp(logcdf) if not implemented.

Author: maresb

pymc/logprob/abstract.py

@@ -40,6 +40,8 @@
 from collections.abc import Sequence
 from functools import singledispatch
 
+import pytensor.tensor as pt
+
 from pytensor.graph import Apply, Op, Variable
 from pytensor.graph.utils import MetaType
 from pytensor.tensor import TensorVariable
@@ -129,8 +131,17 @@ def _logccdf(
 
 
 def _logccdf_helper(rv, value, **kwargs):
-    """Helper that calls `_logccdf` dispatcher."""
-    logccdf = _logccdf(rv.owner.op, value, *rv.owner.inputs, name=rv.name, **kwargs)
+    """Helper that calls `_logccdf` dispatcher with fallback to log1mexp(logcdf).
+
+    If a numerically stable `_logccdf` implementation is registered for the
+    distribution, it will be used. Otherwise, falls back to computing
+    `log(1 - exp(logcdf))` which may be numerically unstable in the tails.
+    """
+    try:
+        logccdf = _logccdf(rv.owner.op, value, *rv.owner.inputs, name=rv.name, **kwargs)
+    except NotImplementedError:
+        logcdf = _logcdf_helper(rv, value, **kwargs)
+        logccdf = pt.log1mexp(logcdf)
 
     if rv.name:
         logccdf.name = f"{rv.name}_logccdf"

pymc/logprob/binary.py

@@ -25,7 +25,7 @@
 
 from pymc.logprob.abstract import (
     MeasurableElemwise,
-    _logccdf,
+    _logccdf_helper,
     _logcdf_helper,
     _logprob,
     _logprob_helper,
@@ -96,12 +96,7 @@ def comparison_logprob(op, values, base_rv, operand, **kwargs):
     base_rv_op = base_rv.owner.op
 
     logcdf = _logcdf_helper(base_rv, operand, **kwargs)
-    # Try to use a numerically stable logccdf if available, otherwise fall back
-    # to computing log(1 - exp(logcdf)) which can be unstable in the tails
-    try:
-        logccdf = _logccdf(base_rv_op, operand, *base_rv.owner.inputs, **kwargs)
-    except NotImplementedError:
-        logccdf = pt.log1mexp(logcdf)
+    logccdf = _logccdf_helper(base_rv, operand, **kwargs)
 
     condn_exp = pt.eq(value, np.array(True))
 

pymc/logprob/censoring.py

@@ -47,7 +47,7 @@
 from pytensor.tensor.math import ceil, clip, floor, round_half_to_even
 from pytensor.tensor.variable import TensorConstant
 
-from pymc.logprob.abstract import MeasurableElemwise, _logccdf, _logcdf, _logprob
+from pymc.logprob.abstract import MeasurableElemwise, _logccdf_helper, _logcdf, _logprob
 from pymc.logprob.rewriting import measurable_ir_rewrites_db
 from pymc.logprob.utils import CheckParameterValue, filter_measurable_variables
 
@@ -119,12 +119,8 @@ def clip_logprob(op, values, base_rv, lower_bound, upper_bound, **kwargs):
     if not (isinstance(upper_bound, TensorConstant) and np.all(np.isinf(upper_bound.value))):
         is_upper_bounded = True
 
-        # Try to use a numerically stable logccdf if available, otherwise fall back
-        # to computing log(1 - exp(logcdf)) which can be unstable in the tails
-        try:
-            logccdf = _logccdf(base_rv_op, value, *base_rv_inputs, **kwargs)
-        except NotImplementedError:
-            logccdf = pt.log1mexp(logcdf)
+        # Use numerically stable logccdf (falls back to log1mexp if not available)
+        logccdf = _logccdf_helper(base_rv, value, **kwargs)
 
         # For right clipped discrete RVs, we need to add an extra term
         # corresponding to the pmf at the upper bound

pymc/logprob/transforms.py

@@ -252,12 +252,8 @@ def measurable_transform_logcdf(op: MeasurableTransform, value, *inputs, **kwarg
         # For discrete distributions, use the logcdf at the previous value
         logccdf = pt.log1mexp(_logcdf_helper(measurable_input, backward_value - 1))
     else:
-        # Try to use a numerically stable logccdf if available, otherwise fall back
-        # to computing log(1 - exp(logcdf)) which can be unstable in the tails
-        try:
-            logccdf = _logccdf_helper(measurable_input, backward_value)
-        except NotImplementedError:
-            logccdf = pt.log1mexp(logcdf)
+        # Use numerically stable logccdf (falls back to log1mexp if not available)
+        logccdf = _logccdf_helper(measurable_input, backward_value)
 
     if isinstance(op.scalar_op, MONOTONICALLY_INCREASING_OPS):
         pass