Return separate logp terms when logpt is called with sum==False

Author: ricardoV94

pymc/distributions/logprob.py

@@ -15,7 +15,7 @@
 
 from collections.abc import Mapping
 from functools import singledispatch
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Union
 
 import aesara.tensor as at
 import numpy as np
@@ -119,15 +119,15 @@ def _get_scaling(total_size, shape, ndim):
 
 
 def logpt(
-    var: TensorVariable,
+    var: Union[TensorVariable, List[TensorVariable]],
     rv_values: Optional[Union[TensorVariable, Dict[TensorVariable, TensorVariable]]] = None,
     *,
     jacobian: bool = True,
     scaling: bool = True,
     transformed: bool = True,
     sum: bool = True,
     **kwargs,
-) -> TensorVariable:
+) -> Union[TensorVariable, List[TensorVariable]]:
     """Create a measure-space (i.e. log-likelihood) graph for a random variable
     or a list of random variables at a given point.
 
@@ -154,7 +154,7 @@ def logpt(
     transformed
         Apply transforms.
     sum
-        Sum the log-likelihood.
+        Sum the log-likelihood or return each term as a separate list item.
 
     """
     # TODO: In future when we drop support for tag.value_var most of the following
@@ -241,7 +241,13 @@ def logpt(
     if sum:
         logp_var = at.sum([at.sum(factor) for factor in logp_var_dict.values()])
     else:
-        logp_var = at.add(*logp_var_dict.values())
+        logp_var = list(logp_var_dict.values())
+        # TODO: deprecate special behavior when only one variable is requested and
+        #  always return a list. This is here for backwards compatibility as logpt
+        #  started as a replacement to factor.logpt, but it should now be considered an
+        #  internal function reached only via model.logp* methods.
+        if len(logp_var) == 1:
+            logp_var = logp_var[0]
 
     return logp_var
 

pymc/tests/test_logprob.py

@@ -144,7 +144,8 @@ def test_logpt_subtensor():
     I_value_var = I_rv.type()
     I_value_var.name = "I_value"
 
-    A_idx_logp = logpt(A_idx, {A_idx: A_idx_value_var, I_rv: I_value_var}, sum=False)
+    A_idx_logps = logpt(A_idx, {A_idx: A_idx_value_var, I_rv: I_value_var}, sum=False)
+    A_idx_logp = at.add(*A_idx_logps)
 
     logp_vals_fn = aesara.function([A_idx_value_var, I_value_var], A_idx_logp)