Separate logp-related dispatch helpers from user facing functions

Author: ricardoV94

pymc/distributions/continuous.py

@@ -56,7 +56,7 @@
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.var import TensorConstant
 
-from pymc.logprob.abstract import _logprob, logcdf, logprob
+from pymc.logprob.abstract import _logcdf_helper, _logprob_helper
 
 try:
     from polyagamma import polyagamma_cdf, polyagamma_pdf, random_polyagamma
@@ -722,7 +722,7 @@ def logp(value, mu, sigma, lower, upper):
         else:
             norm = 0.0
 
-        logp = _logprob(normal, (value,), None, None, None, mu, sigma) - norm
+        logp = _logprob_helper(Normal.dist(mu, sigma), value) - norm
 
         if is_lower_bounded:
             logp = pt.switch(value < lower, -np.inf, logp)
@@ -2033,7 +2033,7 @@ def moment(rv, size, loc, beta):
         return beta
 
     def logp(value, loc, beta):
-        res = pt.log(2) + logprob(Cauchy.dist(loc, beta), value)
+        res = pt.log(2) + _logprob_helper(Cauchy.dist(loc, beta), value)
         res = pt.switch(pt.ge(value, loc), res, -np.inf)
         return check_parameters(
             res,
@@ -2342,10 +2342,10 @@ def moment(rv, size, nu):
         return moment
 
     def logp(value, nu):
-        return logprob(Gamma.dist(alpha=nu / 2, beta=0.5), value)
+        return _logprob_helper(Gamma.dist(alpha=nu / 2, beta=0.5), value)
 
     def logcdf(value, nu):
-        return logcdf(Gamma.dist(alpha=nu / 2, beta=0.5), value)
+        return _logcdf_helper(Gamma.dist(alpha=nu / 2, beta=0.5), value)
 
 
 # TODO: Remove this once logp for multiplication is working!

pymc/distributions/mixture.py

@@ -32,8 +32,7 @@
 )
 from pymc.distributions.shape_utils import _change_dist_size, change_dist_size
 from pymc.distributions.transforms import _default_transform
-from pymc.logprob.abstract import _logcdf, _logprob, logcdf
-from pymc.logprob.basic import logp
+from pymc.logprob.abstract import _logcdf, _logcdf_helper, _logprob, _logprob_helper
 from pymc.logprob.transforms import IntervalTransform
 from pymc.logprob.utils import ignore_logprob
 from pymc.util import check_dist_not_registered
@@ -337,10 +336,10 @@ def marginal_mixture_logprob(op, values, rng, weights, *components, **kwargs):
     if len(components) == 1:
         # Need to broadcast value across mixture axis
         mix_axis = -components[0].owner.op.ndim_supp - 1
-        components_logp = logp(components[0], pt.expand_dims(value, mix_axis))
+        components_logp = _logprob_helper(components[0], pt.expand_dims(value, mix_axis))
     else:
         components_logp = pt.stack(
-            [logp(component, value) for component in components],
+            [_logprob_helper(component, value) for component in components],
             axis=-1,
         )
 
@@ -363,10 +362,10 @@ def marginal_mixture_logcdf(op, value, rng, weights, *components, **kwargs):
     if len(components) == 1:
         # Need to broadcast value across mixture axis
         mix_axis = -components[0].owner.op.ndim_supp - 1
-        components_logcdf = logcdf(components[0], pt.expand_dims(value, mix_axis))
+        components_logcdf = _logcdf_helper(components[0], pt.expand_dims(value, mix_axis))
     else:
         components_logcdf = pt.stack(
-            [logcdf(component, value) for component in components],
+            [_logcdf_helper(component, value) for component in components],
             axis=-1,
         )
 

pymc/distributions/truncated.py

@@ -37,7 +37,8 @@
 from pymc.distributions.shape_utils import _change_dist_size, change_dist_size, to_tuple
 from pymc.distributions.transforms import _default_transform
 from pymc.exceptions import TruncationError
-from pymc.logprob.abstract import MeasurableVariable, _logcdf, _logprob, icdf, logcdf
+from pymc.logprob.abstract import MeasurableVariable, _logcdf, _logprob
+from pymc.logprob.basic import icdf, logcdf
 from pymc.math import logdiffexp
 from pymc.util import check_dist_not_registered
 

pymc/logprob/__init__.py

@@ -34,9 +34,7 @@
 #   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #   SOFTWARE.
 
-from pymc.logprob.abstract import logprob, logcdf  # isort: split
-
-from pymc.logprob.basic import factorized_joint_logprob, joint_logp, logp
+from pymc.logprob.basic import factorized_joint_logprob, icdf, joint_logp, logcdf, logp
 
 # isort: off
 # Add rewrites to the DBs
@@ -49,4 +47,8 @@
 
 # isort: on
 
-__all__ = ("logp", "logcdf")
+__all__ = (
+    "logp",
+    "logcdf",
+    "icdf",
+)

pymc/logprob/abstract.py

@@ -48,35 +48,6 @@
 from pytensor.tensor.random.op import RandomVariable
 
 
-def logprob(rv_var, *rv_values, **kwargs):
-    """Create a graph for the log-probability of a ``RandomVariable``."""
-    logprob = _logprob(rv_var.owner.op, rv_values, *rv_var.owner.inputs, **kwargs)
-
-    for rv_var in rv_values:
-        if rv_var.name:
-            logprob.name = f"{rv_var.name}_logprob"
-
-    return logprob
-
-
-def logcdf(rv_var, rv_value, **kwargs):
-    """Create a graph for the logcdf of a ``RandomVariable``."""
-    logcdf = _logcdf(rv_var.owner.op, rv_value, *rv_var.owner.inputs, name=rv_var.name, **kwargs)
-
-    if rv_var.name:
-        logcdf.name = f"{rv_var.name}_logcdf"
-
-    return logcdf
-
-
-def icdf(rv, value, **kwargs):
-    """Create a graph for the inverse CDF of a `RandomVariable`."""
-    rv_icdf = _icdf(rv.owner.op, value, *rv.owner.inputs, **kwargs)
-    if rv.name:
-        rv_icdf.name = f"{rv.name}_icdf"
-    return rv_icdf
-
-
 @singledispatch
 def _logprob(
     op: Op,
@@ -94,6 +65,18 @@ def _logprob(
     raise NotImplementedError(f"Logprob method not implemented for {op}")
 
 
+def _logprob_helper(rv, *values, **kwargs):
+    """Helper that calls `_logprob` dispatcher."""
+    logprob = _logprob(rv.owner.op, values, *rv.owner.inputs, **kwargs)
+
+    for rv in values:
+        if rv.name:
+            logprob.name = f"{rv.name}_logprob"
+            break
+
+    return logprob
+
+
 @singledispatch
 def _logcdf(
     op: Op,
@@ -107,7 +90,17 @@ def _logcdf(
     of ``RandomVariable``.  If you want to implement new logcdf graphs
     for a ``RandomVariable``, register a new function on this dispatcher.
     """
-    raise NotImplementedError(f"Logcdf method not implemented for {op}")
+    raise NotImplementedError(f"LogCDF method not implemented for {op}")
+
+
+def _logcdf_helper(rv, value, **kwargs):
+    """Helper that calls `_logcdf` dispatcher."""
+    logcdf = _logcdf(rv.owner.op, value, *rv.owner.inputs, name=rv.name, **kwargs)
+
+    if rv.name:
+        logcdf.name = f"{rv.name}_logcdf"
+
+    return logcdf
 
 
 @singledispatch
@@ -122,7 +115,17 @@ def _icdf(
     This function dispatches on the type of `op`, which should be a subclass
     of `RandomVariable`.
     """
-    raise NotImplementedError(f"icdf not implemented for {op}")
+    raise NotImplementedError(f"Inverse CDF method not implemented for {op}")
+
+
+def _icdf_helper(rv, value, **kwargs):
+    """Helper that calls `_icdf` dispatcher."""
+    rv_icdf = _icdf(rv.owner.op, value, *rv.owner.inputs, **kwargs)
+
+    if rv.name:
+        rv_icdf.name = f"{rv.name}_icdf"
+
+    return rv_icdf
 
 
 class MeasurableVariable(abc.ABC):

pymc/logprob/basic.py

@@ -39,33 +39,54 @@
 from collections import deque
 from typing import Dict, List, Optional, Sequence, Union
 
+import numpy as np
 import pytensor
 import pytensor.tensor as pt
 
 from pytensor import config
-from pytensor.graph.basic import graph_inputs, io_toposort
+from pytensor.graph.basic import Variable, graph_inputs, io_toposort
 from pytensor.graph.op import compute_test_value
 from pytensor.graph.rewriting.basic import GraphRewriter, NodeRewriter
 from pytensor.tensor.random.op import RandomVariable
 from pytensor.tensor.var import TensorVariable
-
-from pymc.logprob.abstract import _logprob, get_measurable_outputs
-from pymc.logprob.abstract import logprob as logp_logprob
+from typing_extensions import TypeAlias
+
+from pymc.logprob.abstract import (
+    _icdf_helper,
+    _logcdf_helper,
+    _logprob,
+    _logprob_helper,
+    get_measurable_outputs,
+)
 from pymc.logprob.rewriting import construct_ir_fgraph
 from pymc.logprob.transforms import RVTransform, TransformValuesRewrite
 from pymc.logprob.utils import rvs_to_value_vars
 
+TensorLike: TypeAlias = Union[Variable, float, np.ndarray]
+
 
-def logp(rv: TensorVariable, value: TensorVariable, **kwargs) -> TensorVariable:
+def logp(rv: TensorVariable, value: TensorLike, **kwargs) -> TensorVariable:
     """Return the log-probability graph of a Random Variable"""
 
     value = pt.as_tensor_variable(value, dtype=rv.dtype)
     try:
-        return logp_logprob(rv, value, **kwargs)
+        return _logprob_helper(rv, value, **kwargs)
     except NotImplementedError:
         fgraph, _, _ = construct_ir_fgraph({rv: value})
         [(ir_rv, ir_value)] = fgraph.preserve_rv_mappings.rv_values.items()
-        return logp_logprob(ir_rv, ir_value, **kwargs)
+        return _logprob_helper(ir_rv, ir_value, **kwargs)
+
+
+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)
+
+
+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)
 
 
 def factorized_joint_logprob(

pymc/logprob/cumsum.py

@@ -41,7 +41,7 @@
 from pytensor.graph.rewriting.basic import node_rewriter
 from pytensor.tensor.extra_ops import CumOp
 
-from pymc.logprob.abstract import MeasurableVariable, _logprob, logprob
+from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
 from pymc.logprob.utils import ignore_logprob
 
@@ -72,7 +72,7 @@ def logprob_cumsum(op, values, base_rv, **kwargs):
         axis=op.axis,
     )
 
-    cumsum_logp = logprob(base_rv, value_diff)
+    cumsum_logp = _logprob_helper(base_rv, value_diff)
 
     return cumsum_logp
 

pymc/logprob/mixture.py

@@ -69,7 +69,7 @@
 from pytensor.tensor.type_other import NoneConst, NoneTypeT, SliceConstant, SliceType
 from pytensor.tensor.var import TensorVariable
 
-from pymc.logprob.abstract import MeasurableVariable, _logprob, logprob
+from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
 from pymc.logprob.rewriting import (
     local_lift_DiracDelta,
     logprob_rewrites_db,
@@ -445,7 +445,7 @@ def logprob_MixtureRV(
             # this intentional one-off?
             rv_m = rv_pull_down(rv[m_indices] if m_indices else rv)
             val_m = value[idx_m_on_axis]
-            logp_m = logprob(rv_m, val_m)
+            logp_m = _logprob_helper(rv_m, val_m)
             logp_val = pt.set_subtensor(logp_val[idx_m_on_axis], logp_m)
 
     else:
@@ -463,7 +463,7 @@ def logprob_MixtureRV(
 
         logp_val = 0.0
         for i, comp_rv in enumerate(comp_rvs):
-            comp_logp = logprob(comp_rv, value)
+            comp_logp = _logprob_helper(comp_rv, value)
             if join_axis_val is not None:
                 comp_logp = pt.squeeze(comp_logp, axis=join_axis_val)
             logp_val += ifelse(
@@ -540,10 +540,10 @@ def logprob_ifelse(op, values, if_var, *base_rvs, **kwargs):
     rvs_to_values_else = {else_rv: value for else_rv, value in zip(base_rvs[len(values) :], values)}
 
     logps_then = [
-        logprob(rv_then, value, **kwargs) for rv_then, value in rvs_to_values_then.items()
+        _logprob_helper(rv_then, value, **kwargs) for rv_then, value in rvs_to_values_then.items()
     ]
     logps_else = [
-        logprob(rv_else, value, **kwargs) for rv_else, value in rvs_to_values_else.items()
+        _logprob_helper(rv_else, value, **kwargs) for rv_else, value in rvs_to_values_else.items()
     ]
 
     # If the multiple variables depend on each other, we have to replace them

pymc/logprob/tensor.py

@@ -50,7 +50,7 @@
     local_rv_size_lift,
 )
 
-from pymc.logprob.abstract import MeasurableVariable, _logprob, logprob
+from pymc.logprob.abstract import MeasurableVariable, _logprob, _logprob_helper
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
 from pymc.logprob.utils import ignore_logprob, ignore_logprob_multiple_vars
 
@@ -137,7 +137,7 @@ def logprob_make_vector(op, values, *base_rvs, **kwargs):
         base_rv.name = f"base_rv[{i}]"
         value.name = f"value[{i}]"
 
-    logps = [logprob(base_rv, value) for base_rv, value in base_rvs_to_values.items()]
+    logps = [_logprob_helper(base_rv, value) for base_rv, value in base_rvs_to_values.items()]
 
     # If the stacked variables depend on each other, we have to replace them by the respective values
     logps = replace_rvs_by_values(logps, rvs_to_values=base_rvs_to_values)
@@ -174,7 +174,8 @@ def logprob_join(op, values, axis, *base_rvs, **kwargs):
 
     base_rvs_to_split_values = {base_rv: value for base_rv, value in zip(base_rvs, split_values)}
     logps = [
-        logprob(base_var, split_value) for base_var, split_value in base_rvs_to_split_values.items()
+        _logprob_helper(base_var, split_value)
+        for base_var, split_value in base_rvs_to_split_values.items()
     ]
 
     if len({logp.ndim for logp in logps}) != 1:
@@ -271,7 +272,7 @@ def logprob_dimshuffle(op, values, base_var, **kwargs):
     undo_ds = [original_shuffle.index(i) for i in range(len(original_shuffle))]
     value = value.dimshuffle(undo_ds)
 
-    raw_logp = logprob(base_var, value)
+    raw_logp = _logprob_helper(base_var, value)
 
     # Re-apply original dimshuffle, ignoring any support dimensions consumed by
     # the logprob function. This assumes that support dimensions are always in

pymc/logprob/transforms.py

@@ -78,7 +78,7 @@
     MeasurableVariable,
     _get_measurable_outputs,
     _logprob,
-    logprob,
+    _logprob_helper,
 )
 from pymc.logprob.rewriting import PreserveRVMappings, measurable_ir_rewrites_db
 from pymc.logprob.utils import ignore_logprob, walk_model
@@ -369,10 +369,13 @@ def measurable_transform_logprob(op: MeasurableTransform, values, *inputs, **kwa
     # Some transformations, like squaring may produce multiple backward values
     if isinstance(backward_value, tuple):
         input_logprob = pt.logaddexp(
-            *(logprob(measurable_input, backward_val, **kwargs) for backward_val in backward_value)
+            *(
+                _logprob_helper(measurable_input, backward_val, **kwargs)
+                for backward_val in backward_value
+            )
         )
     else:
-        input_logprob = logprob(measurable_input, backward_value)
+        input_logprob = _logprob_helper(measurable_input, backward_value)
 
     if input_logprob.ndim < value.ndim:
         # Do we just need to sum the jacobian terms across the support dims?