Remove logcdfpt and simplify logp and logcdf helpers

These methods will only retrieve the RV logp and logcdf terms, and no longer attempt to replace any potential RVs in their arguments.
A ValueError message is raised when the provided value is not compatible with the original RV dimensions

Author: ricardoV94

pymc/distributions/__init__.py

@@ -13,10 +13,8 @@
 #   limitations under the License.
 
 from pymc.distributions.logprob import (  # isort:skip
-    _logcdf,
     logcdf,
     logp,
-    logcdfpt,
     logp_transform,
     logpt,
     logpt_sum,
@@ -195,6 +193,5 @@
     "logp",
     "logp_transform",
     "logcdf",
-    "_logcdf",
     "logpt_sum",
 ]

pymc/distributions/discrete.py

@@ -14,7 +14,6 @@
 import aesara.tensor as at
 import numpy as np
 
-from aeppl.logprob import _logprob
 from aesara.tensor.random.basic import (
     RandomVariable,
     bernoulli,
@@ -42,7 +41,7 @@
     normal_lcdf,
 )
 from pymc.distributions.distribution import Discrete
-from pymc.distributions.logprob import _logcdf
+from pymc.distributions.logprob import logcdf, logp
 from pymc.distributions.shape_utils import rv_size_is_none
 from pymc.math import sigmoid
 
@@ -754,7 +753,7 @@ def logp(value, n, p):
         )
 
         # Return Poisson when alpha gets very large.
-        return at.switch(at.gt(alpha, 1e10), Poisson.logp(value, mu), negbinom)
+        return at.switch(at.gt(alpha, 1e10), logp(Poisson.dist(mu=mu), value), negbinom)
 
     def logcdf(value, n, p):
         """
@@ -1371,7 +1370,7 @@ def logp(value, psi, theta):
 
         logp_val = at.switch(
             at.gt(value, 0),
-            at.log(psi) + _logprob(poisson, [value], None, None, None, theta),
+            at.log(psi) + logp(Poisson.dist(mu=theta), value),
             at.logaddexp(at.log1p(-psi), at.log(psi) - theta),
         )
 
@@ -1402,7 +1401,7 @@ def logcdf(value, psi, theta):
         return bound(
             at.logaddexp(
                 at.log1p(-psi),
-                at.log(psi) + _logcdf(poisson, value, {}, theta),
+                at.log(psi) + logcdf(Poisson.dist(mu=theta), value),
             ),
             0 <= value,
             0 <= psi,
@@ -1510,7 +1509,7 @@ def logp(value, psi, n, p):
 
         logp_val = at.switch(
             at.gt(value, 0),
-            at.log(psi) + _logprob(binomial, [value], None, None, None, n, p),
+            at.log(psi) + logp(Binomial.dist(n=n, p=p), value),
             at.logaddexp(at.log1p(-psi), at.log(psi) + n * at.log1p(-p)),
         )
 
@@ -1543,7 +1542,7 @@ def logcdf(value, psi, n, p):
         return bound(
             at.logaddexp(
                 at.log1p(-psi),
-                at.log(psi) + _logcdf(binomial, value, {}, n, p),
+                at.log(psi) + logcdf(Binomial.dist(n=n, p=p), value),
             ),
             0 <= value,
             value <= n,
@@ -1669,7 +1668,7 @@ def logp(value, psi, n, p):
         return bound(
             at.switch(
                 at.gt(value, 0),
-                at.log(psi) + _logprob(nbinom, [value], None, None, None, n, p),
+                at.log(psi) + logp(NegativeBinomial.dist(n=n, p=p), value),
                 at.logaddexp(at.log1p(-psi), at.log(psi) + n * at.log(p)),
             ),
             0 <= value,
@@ -1696,7 +1695,9 @@ def logcdf(value, psi, n, p):
         TensorVariable
         """
         return bound(
-            at.logaddexp(at.log1p(-psi), at.log(psi) + _logcdf(nbinom, value, {}, n, p)),
+            at.logaddexp(
+                at.log1p(-psi), at.log(psi) + logcdf(NegativeBinomial.dist(n=n, p=p), value)
+            ),
             0 <= value,
             0 <= psi,
             psi <= 1,

pymc/distributions/distribution.py

@@ -23,14 +23,13 @@
 
 import aesara
 
-from aeppl.logprob import _logprob
+from aeppl.logprob import _logcdf, _logprob
 from aesara.tensor.basic import as_tensor_variable
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.var import RandomStateSharedVariable
 from aesara.tensor.var import TensorVariable
 
 from pymc.aesaraf import change_rv_size
-from pymc.distributions import _logcdf
 from pymc.distributions.shape_utils import (
     Dims,
     Shape,
@@ -98,18 +97,18 @@ def _random(*args, **kwargs):
             if class_logp:
 
                 @_logprob.register(rv_type)
-                def logp(op, value_var_list, *dist_params, **kwargs):
-                    _dist_params = dist_params[3:]
-                    value_var = value_var_list[0]
-                    return class_logp(value_var, *_dist_params)
+                def logp(op, values, *dist_params, **kwargs):
+                    dist_params = dist_params[3:]
+                    (value,) = values
+                    return class_logp(value, *dist_params)
 
             class_logcdf = clsdict.get("logcdf")
             if class_logcdf:
 
                 @_logcdf.register(rv_type)
-                def logcdf(op, var, rvs_to_values, *dist_params, **kwargs):
-                    value_var = rvs_to_values.get(var, var)
-                    return class_logcdf(value_var, *dist_params, **kwargs)
+                def logcdf(op, value, *dist_params, **kwargs):
+                    dist_params = dist_params[3:]
+                    return class_logcdf(value, *dist_params)
 
             class_initval = clsdict.get("get_moment")
             if class_initval:

pymc/distributions/logprob.py

@@ -21,6 +21,8 @@
 import numpy as np
 
 from aeppl import factorized_joint_logprob
+from aeppl.logprob import logcdf as logcdf_aeppl
+from aeppl.logprob import logprob as logp_aeppl
 from aeppl.transforms import TransformValuesOpt
 from aesara import config
 from aesara.graph.basic import graph_inputs, io_toposort
@@ -36,7 +38,7 @@
 )
 from aesara.tensor.var import TensorVariable
 
-from pymc.aesaraf import extract_rv_and_value_vars, floatX, rvs_to_value_vars
+from pymc.aesaraf import floatX
 
 
 @singledispatch
@@ -260,128 +262,18 @@ def logpt(
     return logp_var
 
 
-def logcdfpt(
-    var: TensorVariable,
-    rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None,
-    *,
-    scaling: bool = True,
-    sum: bool = True,
-    **kwargs,
-) -> TensorVariable:
-    """Create a measure-space (i.e. log-cdf) graph for a random variable at a given point.
-
-    Parameters
-    ==========
-    var
-        The `RandomVariable` output that determines the log-likelihood graph.
-    rv_values
-        A variable, or ``dict`` of variables, that represents the value of
-        `var` in its log-likelihood.  If no `rv_value` is provided,
-        ``var.tag.value_var`` will be checked and, when available, used.
-    jacobian
-        Whether or not to include the Jacobian term.
-    scaling
-        A scaling term to apply to the generated log-likelihood graph.
-    transformed
-        Apply transforms.
-    sum
-        Sum the log-likelihood.
-
-    """
-    if not isinstance(rv_values, Mapping):
-        rv_values = {var: rv_values} if rv_values is not None else {}
-
-    rv_var, rv_value_var = extract_rv_and_value_vars(var)
+def logp(rv, value):
+    """Return the log-probability graph of a Random Variable"""
 
-    rv_value = rv_values.get(rv_var, rv_value_var)
+    value = at.as_tensor_variable(value, dtype=rv.dtype)
+    return logp_aeppl(rv, value)
 
-    if rv_var is not None and rv_value is None:
-        raise ValueError(f"No value variable specified or associated with {rv_var}")
 
-    if rv_value is not None:
-        rv_value = at.as_tensor(rv_value)
-
-        if rv_var is not None:
-            # Make sure that the value is compatible with the random variable
-            rv_value = rv_var.type.filter_variable(rv_value.astype(rv_var.dtype))
-
-        if rv_value_var is None:
-            rv_value_var = rv_value
-
-    rv_node = rv_var.owner
-
-    rng, size, dtype, *dist_params = rv_node.inputs
-
-    # Here, we plug the actual random variable into the log-likelihood graph,
-    # because we want a log-likelihood graph that only contains
-    # random variables.  This is important, because a random variable's
-    # parameters can contain random variables themselves.
-    # Ultimately, with a graph containing only random variables and
-    # "deterministics", we can simply replace all the random variables with
-    # their value variables and be done.
-    tmp_rv_values = rv_values.copy()
-    tmp_rv_values[rv_var] = rv_var
-
-    logp_var = _logcdf(rv_node.op, rv_var, tmp_rv_values, *dist_params, **kwargs)
-
-    transform = getattr(rv_value_var.tag, "transform", None) if rv_value_var else None
-
-    # Replace random variables with their value variables
-    replacements = rv_values.copy()
-    replacements.update({rv_var: rv_value, rv_value_var: rv_value})
-
-    (logp_var,), _ = rvs_to_value_vars(
-        (logp_var,),
-        apply_transforms=False,
-        initial_replacements=replacements,
-    )
-
-    if sum:
-        logp_var = at.sum(logp_var)
-
-    if scaling:
-        logp_var *= _get_scaling(
-            getattr(rv_var.tag, "total_size", None), rv_value.shape, rv_value.ndim
-        )
-
-    # Recompute test values for the changes introduced by the replacements
-    # above.
-    if config.compute_test_value != "off":
-        for node in io_toposort(graph_inputs((logp_var,)), (logp_var,)):
-            compute_test_value(node)
-
-    if rv_var.name is not None:
-        logp_var.name = f"__logp_{rv_var.name}"
-
-    return logp_var
-
-
-def logp(var, rv_values, **kwargs):
-    """Create a log-probability graph."""
-
-    # Attach the value_var to the tag of var when it does not have one
-    if not hasattr(var.tag, "value_var"):
-        if isinstance(rv_values, Mapping):
-            value_var = rv_values[var]
-        else:
-            value_var = rv_values
-        var.tag.value_var = at.as_tensor_variable(value_var, dtype=var.dtype)
-
-    return logpt(var, rv_values, **kwargs)
-
-
-def logcdf(var, rv_values, **kwargs):
-    """Create a log-CDF graph."""
-
-    # Attach the value_var to the tag of var when it does not have one
-    if not hasattr(var.tag, "value_var"):
-        if isinstance(rv_values, Mapping):
-            value_var = rv_values[var]
-        else:
-            value_var = rv_values
-        var.tag.value_var = at.as_tensor_variable(value_var, dtype=var.dtype)
+def logcdf(rv, value):
+    """Return the log-cdf graph of a Random Variable"""
 
-    return logcdfpt(var, rv_values, **kwargs)
+    value = at.as_tensor_variable(value, dtype=rv.dtype)
+    return logcdf_aeppl(rv, value)
 
 
 @singledispatch
@@ -402,4 +294,5 @@ def logpt_sum(*args, **kwargs):
     Subclasses can use this to improve the speed of logp evaluations
     if only the sum of the logp values is needed.
     """
+    # TODO: Deprecate this
     return logpt(*args, sum=True, **kwargs)

pymc/distributions/multivariate.py

@@ -2102,7 +2102,7 @@ def logp(value, mu, W, alpha, tau):
         TensorVariable
         """
 
-        sparse = isinstance(W, aesara.sparse.SparseConstant)
+        sparse = isinstance(W, (aesara.sparse.SparseConstant, aesara.sparse.SparseVariable))
 
         if sparse:
             D = sp_sum(W, axis=0)