Remove joint_logprob function from tests.logprob.utils

Author: shreyas3156

tests/logprob/test_censoring.py

@@ -41,10 +41,10 @@
 import scipy as sp
 import scipy.stats as st
 
+from pymc import logp
 from pymc.logprob import factorized_joint_logprob
 from pymc.logprob.transforms import LogTransform, TransformValuesRewrite
 from pymc.testing import assert_no_rvs
-from tests.logprob.utils import joint_logprob
 
 
 @pytensor.config.change_flags(compute_test_value="raise")
@@ -55,10 +55,10 @@ def test_continuous_rv_clip():
     cens_x_vv = cens_x_rv.clone()
     cens_x_vv.tag.test_value = 0
 
-    logp = joint_logprob({cens_x_rv: cens_x_vv})
-    assert_no_rvs(logp)
+    logprob = pt.sum(logp(cens_x_rv, cens_x_vv))
+    assert_no_rvs(logprob)
 
-    logp_fn = pytensor.function([cens_x_vv], logp)
+    logp_fn = pytensor.function([cens_x_vv], logprob)
     ref_scipy = st.norm(0.5, 1)
 
     assert logp_fn(-3) == -np.inf
@@ -75,10 +75,10 @@ def test_discrete_rv_clip():
 
     cens_x_vv = cens_x_rv.clone()
 
-    logp = joint_logprob({cens_x_rv: cens_x_vv})
-    assert_no_rvs(logp)
+    logprob = pt.sum(logp(cens_x_rv, cens_x_vv))
+    assert_no_rvs(logprob)
 
-    logp_fn = pytensor.function([cens_x_vv], logp)
+    logp_fn = pytensor.function([cens_x_vv], logprob)
     ref_scipy = st.poisson(2)
 
     assert logp_fn(0) == -np.inf
@@ -97,8 +97,8 @@ def test_one_sided_clip():
     lb_cens_x_vv = lb_cens_x_rv.clone()
     ub_cens_x_vv = ub_cens_x_rv.clone()
 
-    lb_logp = joint_logprob({lb_cens_x_rv: lb_cens_x_vv})
-    ub_logp = joint_logprob({ub_cens_x_rv: ub_cens_x_vv})
+    lb_logp = pt.sum(logp(lb_cens_x_rv, lb_cens_x_vv))
+    ub_logp = pt.sum(logp(ub_cens_x_rv, ub_cens_x_vv))
     assert_no_rvs(lb_logp)
     assert_no_rvs(ub_logp)
 
@@ -117,10 +117,10 @@ def test_useless_clip():
 
     cens_x_vv = cens_x_rv.clone()
 
-    logp = joint_logprob({cens_x_rv: cens_x_vv}, sum=False)
-    assert_no_rvs(logp)
+    logprob = logp(cens_x_rv, cens_x_vv)
+    assert_no_rvs(logprob)
 
-    logp_fn = pytensor.function([cens_x_vv], logp)
+    logp_fn = pytensor.function([cens_x_vv], logprob)
     ref_scipy = st.norm(0.5, 1)
 
     np.testing.assert_allclose(logp_fn([-2, 0, 2]), ref_scipy.logpdf([-2, 0, 2]))
@@ -133,10 +133,12 @@ def test_random_clip():
 
     lb_vv = lb_rv.clone()
     cens_x_vv = cens_x_rv.clone()
-    logp = joint_logprob({cens_x_rv: cens_x_vv, lb_rv: lb_vv}, sum=False)
-    assert_no_rvs(logp)
+    logp = factorized_joint_logprob({cens_x_rv: cens_x_vv, lb_rv: lb_vv})
+    logp_combined = pt.add(*logp.values())
 
-    logp_fn = pytensor.function([lb_vv, cens_x_vv], logp)
+    assert_no_rvs(logp_combined)
+
+    logp_fn = pytensor.function([lb_vv, cens_x_vv], logp_combined)
     res = logp_fn([0, -1], [-1, -1])
     assert res[0] == -np.inf
     assert res[1] != -np.inf
@@ -150,8 +152,10 @@ def test_broadcasted_clip_constant():
     lb_vv = lb_rv.clone()
     cens_x_vv = cens_x_rv.clone()
 
-    logp = joint_logprob({cens_x_rv: cens_x_vv, lb_rv: lb_vv})
-    assert_no_rvs(logp)
+    logp = factorized_joint_logprob({cens_x_rv: cens_x_vv, lb_rv: lb_vv})
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
+
+    assert_no_rvs(logp_combined)
 
 
 def test_broadcasted_clip_random():
@@ -162,8 +166,10 @@ def test_broadcasted_clip_random():
     lb_vv = lb_rv.clone()
     cens_x_vv = cens_x_rv.clone()
 
-    logp = joint_logprob({cens_x_rv: cens_x_vv, lb_rv: lb_vv})
-    assert_no_rvs(logp)
+    logp = factorized_joint_logprob({cens_x_rv: cens_x_vv, lb_rv: lb_vv})
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
+
+    assert_no_rvs(logp_combined)
 
 
 def test_fail_base_and_clip_have_values():
@@ -199,10 +205,11 @@ def test_deterministic_clipping():
 
     x_vv = x_rv.clone()
     y_vv = y_rv.clone()
-    logp = joint_logprob({x_rv: x_vv, y_rv: y_vv})
-    assert_no_rvs(logp)
+    logp = factorized_joint_logprob({x_rv: x_vv, y_rv: y_vv})
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
+    assert_no_rvs(logp_combined)
 
-    logp_fn = pytensor.function([x_vv, y_vv], logp)
+    logp_fn = pytensor.function([x_vv, y_vv], logp_combined)
     assert np.isclose(
         logp_fn(-1, 1),
         st.norm(0, 1).logpdf(-1) + st.norm(0, 1).logpdf(1),
@@ -216,10 +223,11 @@ def test_clip_transform():
     cens_x_vv = cens_x_rv.clone()
 
     transform = TransformValuesRewrite({cens_x_vv: LogTransform()})
-    logp = joint_logprob({cens_x_rv: cens_x_vv}, extra_rewrites=transform)
+    logp = factorized_joint_logprob({cens_x_rv: cens_x_vv}, extra_rewrites=transform)
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
 
     cens_x_vv_testval = -1
-    obs_logp = logp.eval({cens_x_vv: cens_x_vv_testval})
+    obs_logp = logp_combined.eval({cens_x_vv: cens_x_vv_testval})
     exp_logp = sp.stats.norm(0.5, 1).logpdf(np.exp(cens_x_vv_testval)) + cens_x_vv_testval
 
     assert np.isclose(obs_logp, exp_logp)
@@ -236,8 +244,8 @@ def test_rounding(rounding_op):
     xr.name = "xr"
 
     xr_vv = xr.clone()
-    logp = joint_logprob({xr: xr_vv}, sum=False)
-    assert logp is not None
+    logprob = logp(xr, xr_vv)
+    assert logprob is not None
 
     x_sp = st.norm(loc, scale)
     if rounding_op == pt.round:
@@ -250,6 +258,6 @@ def test_rounding(rounding_op):
         raise NotImplementedError()
 
     assert np.allclose(
-        logp.eval({xr_vv: test_value}),
+        logprob.eval({xr_vv: test_value}),
         expected_logp,
     )

tests/logprob/test_composite_logprob.py

@@ -39,10 +39,11 @@
 import pytensor.tensor as pt
 import scipy.stats as st
 
+from pymc import logp
+from pymc.logprob.basic import factorized_joint_logprob
 from pymc.logprob.censoring import MeasurableClip
 from pymc.logprob.rewriting import construct_ir_fgraph
 from pymc.testing import assert_no_rvs
-from tests.logprob.utils import joint_logprob
 
 
 def test_scalar_clipped_mixture():
@@ -61,9 +62,10 @@ def test_scalar_clipped_mixture():
     idxs_vv = idxs.clone()
     idxs_vv.name = "idxs_val"
 
-    logp = joint_logprob({idxs: idxs_vv, mix: mix_vv})
+    logp = factorized_joint_logprob({idxs: idxs_vv, mix: mix_vv})
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
 
-    logp_fn = pytensor.function([idxs_vv, mix_vv], logp)
+    logp_fn = pytensor.function([idxs_vv, mix_vv], logp_combined)
     assert logp_fn(0, 0.4) == -np.inf
     assert np.isclose(logp_fn(0, 0.5), st.norm.logcdf(0.5, 1) + np.log(0.6))
     assert np.isclose(logp_fn(0, 1.3), st.norm.logpdf(1.3, 1) + np.log(0.6))
@@ -98,8 +100,12 @@ def test_nested_scalar_mixtures():
     idxs12_vv = idxs12.clone()
     mix12_vv = mix12.clone()
 
-    logp = joint_logprob({idxs1: idxs1_vv, idxs2: idxs2_vv, idxs12: idxs12_vv, mix12: mix12_vv})
-    logp_fn = pytensor.function([idxs1_vv, idxs2_vv, idxs12_vv, mix12_vv], logp)
+    logp = factorized_joint_logprob(
+        {idxs1: idxs1_vv, idxs2: idxs2_vv, idxs12: idxs12_vv, mix12: mix12_vv}
+    )
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
+
+    logp_fn = pytensor.function([idxs1_vv, idxs2_vv, idxs12_vv, mix12_vv], logp_combined)
 
     expected_mu_logpdf = st.norm.logpdf(0) + np.log(0.5) * 3
     assert np.isclose(logp_fn(0, 0, 0, -50), expected_mu_logpdf)
@@ -144,9 +150,9 @@ def test_shifted_cumsum():
     y.name = "y"
 
     y_vv = y.clone()
-    logp = joint_logprob({y: y_vv})
+    logprob = logp(y, y_vv)
     assert np.isclose(
-        logp.eval({y_vv: np.arange(5) + 1 + 5}),
+        logprob.eval({y_vv: np.arange(5) + 1 + 5}).sum(),
         st.norm.logpdf(1) * 5,
     )
 
@@ -157,8 +163,8 @@ def test_double_log_transform_rv():
     y_rv.name = "y"
 
     y_vv = y_rv.clone()
-    logp = joint_logprob({y_rv: y_vv}, sum=False)
-    logp_fn = pytensor.function([y_vv], logp)
+    logprob = logp(y_rv, y_vv)
+    logp_fn = pytensor.function([y_vv], logprob)
 
     log_log_y_val = np.asarray(0.5)
     log_y_val = np.exp(log_log_y_val)
@@ -178,9 +184,9 @@ def test_affine_transform_rv():
     y_rv.name = "y"
     y_vv = y_rv.clone()
 
-    logp = joint_logprob({y_rv: y_vv}, sum=False)
-    assert_no_rvs(logp)
-    logp_fn = pytensor.function([loc, scale, y_vv], logp)
+    logprob = logp(y_rv, y_vv)
+    assert_no_rvs(logprob)
+    logp_fn = pytensor.function([loc, scale, y_vv], logprob)
 
     loc_test_val = 4.0
     scale_test_val = np.full(rv_size, 0.5)
@@ -200,8 +206,8 @@ def test_affine_log_transform_rv():
 
     y_vv = y_rv.clone()
 
-    logp = joint_logprob({y_rv: y_vv}, sum=False)
-    logp_fn = pytensor.function([a, b, y_vv], logp)
+    logprob = logp(y_rv, y_vv)
+    logp_fn = pytensor.function([a, b, y_vv], logprob)
 
     a_val = -1.5
     b_val = 3.0

tests/logprob/test_cumsum.py

@@ -40,8 +40,9 @@
 import pytest
 import scipy.stats as st
 
+from pymc import logp
+from pymc.logprob.basic import factorized_joint_logprob
 from pymc.testing import assert_no_rvs
-from tests.logprob.utils import joint_logprob
 
 
 @pytest.mark.parametrize(
@@ -59,12 +60,12 @@
 def test_normal_cumsum(size, axis):
     rv = pt.random.normal(0, 1, size=size).cumsum(axis)
     vv = rv.clone()
-    logp = joint_logprob({rv: vv})
-    assert_no_rvs(logp)
+    logprob = logp(rv, vv)
+    assert_no_rvs(logprob)
 
     assert np.isclose(
         st.norm(0, 1).logpdf(np.ones(size)).sum(),
-        logp.eval({vv: np.ones(size).cumsum(axis)}),
+        logprob.eval({vv: np.ones(size).cumsum(axis)}).sum(),
     )
 
 
@@ -83,12 +84,12 @@ def test_normal_cumsum(size, axis):
 def test_bernoulli_cumsum(size, axis):
     rv = pt.random.bernoulli(0.9, size=size).cumsum(axis)
     vv = rv.clone()
-    logp = joint_logprob({rv: vv})
-    assert_no_rvs(logp)
+    logprob = logp(rv, vv)
+    assert_no_rvs(logprob)
 
     assert np.isclose(
         st.bernoulli(0.9).logpmf(np.ones(size)).sum(),
-        logp.eval({vv: np.ones(size, int).cumsum(axis)}),
+        logprob.eval({vv: np.ones(size, int).cumsum(axis)}).sum(),
     )
 
 
@@ -97,7 +98,7 @@ def test_destructive_cumsum_fails():
     x_rv = pt.random.normal(size=(2, 2, 2)).cumsum()
     x_vv = x_rv.clone()
     with pytest.raises(RuntimeError, match="could not be derived"):
-        joint_logprob({x_rv: x_vv})
+        factorized_joint_logprob({x_rv: x_vv})
 
 
 def test_deterministic_cumsum():
@@ -108,11 +109,13 @@ def test_deterministic_cumsum():
 
     x_vv = x_rv.clone()
     y_vv = y_rv.clone()
-    logp = joint_logprob({x_rv: x_vv, y_rv: y_vv})
-    assert_no_rvs(logp)
 
-    logp_fn = pytensor.function([x_vv, y_vv], logp)
+    logp = factorized_joint_logprob({x_rv: x_vv, y_rv: y_vv})
+    logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
+    assert_no_rvs(logp_combined)
+
+    logp_fn = pytensor.function([x_vv, y_vv], logp_combined)
     assert np.isclose(
-        logp_fn(np.ones(5), np.arange(5) + 1),
+        logp_fn(np.ones(5), np.arange(5) + 1).sum(),
         st.norm(1, 1).logpdf(1) * 10,
     )

tests/logprob/test_mixture.py

@@ -52,7 +52,7 @@
 from pymc.logprob.rewriting import construct_ir_fgraph
 from pymc.logprob.utils import dirac_delta
 from pymc.testing import assert_no_rvs
-from tests.logprob.utils import joint_logprob, scipy_logprob
+from tests.logprob.utils import scipy_logprob
 
 
 def test_mixture_basics():
@@ -101,7 +101,7 @@ def create_mix_model(size, axis):
         i_vv = env["i_vv"]
         M_rv = env["M_rv"]
         m_vv = env["m_vv"]
-        joint_logprob({M_rv: m_vv, I_rv: i_vv})
+        factorized_joint_logprob({M_rv: m_vv, I_rv: i_vv})
 
 
 @pytensor.config.change_flags(compute_test_value="warn")
@@ -134,9 +134,10 @@ def test_compute_test_value(op_constructor):
 
     del M_rv.tag.test_value
 
-    M_logp = joint_logprob({M_rv: m_vv, I_rv: i_vv}, sum=False)
+    M_logp = factorized_joint_logprob({M_rv: m_vv, I_rv: i_vv})
+    M_logp_combined = pt.add(*M_logp.values())
 
-    assert isinstance(M_logp.tag.test_value, np.ndarray)
+    assert isinstance(M_logp_combined.tag.test_value, np.ndarray)
 
 
 @pytest.mark.parametrize(
@@ -179,13 +180,14 @@ def test_hetero_mixture_binomial(p_val, size, supported):
     m_vv.name = "m"
 
     if supported:
-        M_logp = joint_logprob({M_rv: m_vv, I_rv: i_vv}, sum=False)
+        M_logp = factorized_joint_logprob({M_rv: m_vv, I_rv: i_vv})
+        M_logp_combined = pt.add(*M_logp.values())
     else:
         with pytest.raises(RuntimeError, match="could not be derived: {m}"):
-            joint_logprob({M_rv: m_vv, I_rv: i_vv}, sum=False)
+            factorized_joint_logprob({M_rv: m_vv, I_rv: i_vv})
         return
 
-    M_logp_fn = pytensor.function([p_at, m_vv, i_vv], M_logp)
+    M_logp_fn = pytensor.function([p_at, m_vv, i_vv], M_logp_combined)
 
     assert_no_rvs(M_logp_fn.maker.fgraph.outputs[0])
 
@@ -936,14 +938,16 @@ def test_switch_mixture():
 
     assert equal_computations(fgraph.outputs, fgraph2.outputs)
 
-    z1_logp = joint_logprob({Z1_rv: z_vv, I_rv: i_vv})
-    z2_logp = joint_logprob({Z2_rv: z_vv, I_rv: i_vv})
+    z1_logp = factorized_joint_logprob({Z1_rv: z_vv, I_rv: i_vv})
+    z2_logp = factorized_joint_logprob({Z2_rv: z_vv, I_rv: i_vv})
+    z1_logp_combined = pt.sum([pt.sum(factor) for factor in z1_logp.values()])
+    z2_logp_combined = pt.sum([pt.sum(factor) for factor in z2_logp.values()])
 
     # below should follow immediately from the equal_computations assertion above
-    assert equal_computations([z1_logp], [z2_logp])
+    assert equal_computations([z1_logp_combined], [z2_logp_combined])
 
-    np.testing.assert_almost_equal(0.69049938, z1_logp.eval({z_vv: -10, i_vv: 0}))
-    np.testing.assert_almost_equal(0.69049938, z2_logp.eval({z_vv: -10, i_vv: 0}))
+    np.testing.assert_almost_equal(0.69049938, z1_logp_combined.eval({z_vv: -10, i_vv: 0}))
+    np.testing.assert_almost_equal(0.69049938, z2_logp_combined.eval({z_vv: -10, i_vv: 0}))
 
 
 def test_ifelse_mixture_one_component():