Make joint_logp always return lists when sum=False

Author: michaelosthege

pymc/distributions/logprob.py

@@ -246,12 +246,6 @@ def joint_logpt(
         logp_var = at.sum([at.sum(factor) for factor in logp_var_dict.values()])
     else:
         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/model.py

@@ -749,11 +749,10 @@ def logpt(
                         f"Requested variable {var} not found among the model variables"
                     )
 
-        rv_logps = []
+        rv_logps: List[TensorVariable] = []
         if rv_values:
             rv_logps = joint_logpt(list(rv_values.keys()), rv_values, sum=False, jacobian=jacobian)
-            if not isinstance(rv_logps, list):
-                rv_logps = [rv_logps]
+            assert isinstance(rv_logps, list)
 
         # Replace random variables by their value variables in potential terms
         potential_logps = []

pymc/tests/test_distributions.py

@@ -2783,19 +2783,19 @@ def test_array_bound(self):
             UpperPoisson = Bound("upper", dist, upper=[np.inf, 10], transform=None)
             BoundedPoisson = Bound("bounded", dist, lower=[1, 2], upper=[9, 10], transform=None)
 
-        first, second = joint_logpt(LowerPoisson, [0, 0], sum=False).eval()
+        first, second = joint_logpt(LowerPoisson, [0, 0], sum=False)[0].eval()
         assert first == -np.inf
         assert second != -np.inf
 
-        first, second = joint_logpt(UpperPoisson, [11, 11], sum=False).eval()
+        first, second = joint_logpt(UpperPoisson, [11, 11], sum=False)[0].eval()
         assert first != -np.inf
         assert second == -np.inf
 
-        first, second = joint_logpt(BoundedPoisson, [1, 1], sum=False).eval()
+        first, second = joint_logpt(BoundedPoisson, [1, 1], sum=False)[0].eval()
         assert first != -np.inf
         assert second == -np.inf
 
-        first, second = joint_logpt(BoundedPoisson, [10, 10], sum=False).eval()
+        first, second = joint_logpt(BoundedPoisson, [10, 10], sum=False)[0].eval()
         assert first == -np.inf
         assert second != -np.inf
 
@@ -3285,7 +3285,7 @@ def logp(value, mu):
     a_val = np.random.normal(loc=mu_val, scale=1, size=to_tuple(size) + (supp_shape,)).astype(
         aesara.config.floatX
     )
-    log_densityt = joint_logpt(a, a.tag.value_var, sum=False)
+    log_densityt = joint_logpt(a, a.tag.value_var, sum=False)[0]
     assert log_densityt.eval(
         {a.tag.value_var: a_val, mu.tag.value_var: mu_val},
     ).shape == to_tuple(size)

pymc/tests/test_logprob.py

@@ -60,7 +60,7 @@ def test_joint_logpt_basic():
 
     b_logp = joint_logpt(b, b_value_var, sum=False)
 
-    res_ancestors = list(walk_model((b_logp,), walk_past_rvs=True))
+    res_ancestors = list(walk_model(b_logp, walk_past_rvs=True))
     res_rv_ancestors = [
         v for v in res_ancestors if v.owner and isinstance(v.owner.op, RandomVariable)
     ]
@@ -104,7 +104,7 @@ def test_joint_logpt_incsubtensor(indices, size):
 
     a_idx_logp = joint_logpt(a_idx, {a_idx: a_value_var}, sum=False)
 
-    logp_vals = a_idx_logp.eval({a_value_var: a_val})
+    logp_vals = a_idx_logp[0].eval({a_value_var: a_val})
 
     # The indices that were set should all have the same log-likelihood values,
     # because the values they were set to correspond to the unique means along

pymc/tests/test_transforms.py

@@ -296,7 +296,7 @@ def check_transform_elementwise_logp(self, model):
         x_val_untransf = at.constant(test_array_untransf).type()
 
         jacob_det = transform.log_jac_det(test_array_transf, *x.owner.inputs)
-        assert joint_logpt(x, sum=False).ndim == x.ndim == jacob_det.ndim
+        assert joint_logpt(x, sum=False)[0].ndim == x.ndim == jacob_det.ndim
 
         v1 = joint_logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf})
         v2 = joint_logpt(x, x_val_untransf, transformed=False).eval(
@@ -319,10 +319,10 @@ def check_vectortransform_elementwise_logp(self, model):
         jacob_det = transform.log_jac_det(test_array_transf, *x.owner.inputs)
         # Original distribution is univariate
         if x.owner.op.ndim_supp == 0:
-            assert joint_logpt(x, sum=False).ndim == x.ndim == (jacob_det.ndim + 1)
+            assert joint_logpt(x, sum=False)[0].ndim == x.ndim == (jacob_det.ndim + 1)
         # Original distribution is multivariate
         else:
-            assert joint_logpt(x, sum=False).ndim == (x.ndim - 1) == jacob_det.ndim
+            assert joint_logpt(x, sum=False)[0].ndim == (x.ndim - 1) == jacob_det.ndim
 
         a = joint_logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf})
         b = joint_logpt(x, x_val_untransf, transformed=False).eval(