Rename basic "joint_logprob" functions to "conditional_logp"

Author: ricardoV94

pymc/logprob/__init__.py

@@ -34,7 +34,13 @@
 #   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 #   SOFTWARE.
 
-from pymc.logprob.basic import factorized_joint_logprob, icdf, joint_logp, logcdf, logp
+from pymc.logprob.basic import (
+    conditional_logp,
+    icdf,
+    logcdf,
+    logp,
+    transformed_conditional_logp,
+)
 
 # isort: off
 # Add rewrites to the DBs

pymc/logprob/basic.py

@@ -163,15 +163,15 @@ def icdf(
 )
 
 
-def factorized_joint_logprob(
+def conditional_logp(
     rv_values: Dict[TensorVariable, TensorVariable],
     warn_missing_rvs: bool = True,
     ir_rewriter: Optional[GraphRewriter] = None,
     extra_rewrites: Optional[Union[GraphRewriter, NodeRewriter]] = None,
     **kwargs,
 ) -> Dict[TensorVariable, TensorVariable]:
-    r"""Create a map between variables and their log-probabilities such that the
-    sum is their joint log-probability.
+    r"""Create a map between variables and conditional log-probabilities
+    such that the sum is their joint log-probability.
 
     The `rv_values` dictionary specifies a joint probability graph defined by
     pairs of random variables and respective measure-space input parameters
@@ -189,20 +189,21 @@ def factorized_joint_logprob(
 
     .. math::
 
-        \sigma^2 \sim& \operatorname{InvGamma}(0.5, 0.5) \\
-        Y \sim& \operatorname{N}(0, \sigma^2)
+        \Sigma^2 \sim& \operatorname{InvGamma}(0.5, 0.5) \\
+        Y \sim& \operatorname{N}(0, \Sigma)
 
     If we create a value variable for ``Y_rv``, i.e. ``y_vv = pt.scalar("y")``,
-    the graph of ``factorized_joint_logprob({Y_rv: y_vv})`` is equivalent to the
-    conditional probability :math:`\log p(Y = y \mid \sigma^2)`, with a stochastic
+    the graph of ``conditional_logp({Y_rv: y_vv})`` is equivalent to the
+    conditional log-probability :math:`\log p(Y = y \mid \Sigma^2)`, with a stochastic
     ``sigma2_rv``. If we specify a value variable for ``sigma2_rv``, i.e.
-    ``s_vv = pt.scalar("s2")``, then ``factorized_joint_logprob({Y_rv: y_vv, sigma2_rv: s_vv})``
-    yields the joint log-probability of the two variables.
+    ``s_vv = pt.scalar("s2")``, then ``conditional_logp({Y_rv: y_vv, sigma2_rv: s_vv})``
+    yields the conditional log-probabilities of the two variables.
+    The sum of the two terms gives their joint log-probability.
 
     .. math::
 
-        \log p(Y = y, \sigma^2 = s) =
-            \log p(Y = y \mid \sigma^2 = s) + \log p(\sigma^2 = s)
+        \log p(Y = y, \Sigma^2 = \sigma^2) =
+            \log p(Y = y \mid \Sigma^2 = \sigma^2) + \log p(\Sigma^2 = \sigma^2)
 
 
     Parameters
@@ -223,7 +224,7 @@ def factorized_joint_logprob(
 
     Returns
     -------
-    A ``dict`` that maps each value variable to the log-probability factor derived
+    A ``dict`` that maps each value variable to the conditional log-probability term derived
     from the respective `RandomVariable`.
 
     """
@@ -309,7 +310,7 @@ def factorized_joint_logprob(
 
             if q_value_var in logprob_vars:
                 raise ValueError(
-                    f"More than one logprob factor was assigned to the value var {q_value_var}"
+                    f"More than one logprob term was assigned to the value var {q_value_var}"
                 )
 
             logprob_vars[q_value_var] = q_logprob_var
@@ -337,16 +338,19 @@ def factorized_joint_logprob(
     return logprob_vars
 
 
-def joint_logp(
+def transformed_conditional_logp(
     rvs: Sequence[TensorVariable],
     *,
     rvs_to_values: Dict[TensorVariable, TensorVariable],
     rvs_to_transforms: Dict[TensorVariable, RVTransform],
     jacobian: bool = True,
     **kwargs,
 ) -> List[TensorVariable]:
-    """Thin wrapper around pymc.logprob.factorized_joint_logprob, extended with Model
-    specific concerns such as transforms, jacobian, and scaling"""
+    """Thin wrapper around conditional_logprob, which creates a value transform rewrite.
+
+    This helper will only return the subset of logprob terms corresponding to `rvs`.
+    All rvs_to_values and rvs_to_transforms mappings are required.
+    """
 
     transform_rewrite = None
     values_to_transforms = {
@@ -359,7 +363,7 @@ def joint_logp(
         transform_rewrite = TransformValuesRewrite(values_to_transforms)  # type: ignore
 
     kwargs.setdefault("warn_missing_rvs", False)
-    temp_logp_terms = factorized_joint_logprob(
+    temp_logp_terms = conditional_logp(
         rvs_to_values,
         extra_rewrites=transform_rewrite,
         use_jacobian=jacobian,
@@ -381,3 +385,21 @@ def joint_logp(
         raise ValueError(RVS_IN_JOINT_LOGP_GRAPH_MSG % rvs_in_logp_expressions)
 
     return logp_terms_list
+
+
+def factorized_joint_logprob(*args, **kwargs):
+    warnings.warn(
+        "`factorized_joint_logprob` was renamed to `conditional_logp`. "
+        "The function will be removed in a future release",
+        FutureWarning,
+    )
+    return conditional_logp(*args, **kwargs)
+
+
+def joint_logp(*args, **kwargs):
+    warnings.warn(
+        "`joint_logp` was renamed to `transformed_conditional_logp`. "
+        "The function will be removed in a future release",
+        FutureWarning,
+    )
+    return transformed_conditional_logp(*args, **kwargs)

pymc/logprob/scan.py

@@ -54,7 +54,7 @@
 from pytensor.updates import OrderedUpdates
 
 from pymc.logprob.abstract import MeasurableVariable, _logprob
-from pymc.logprob.basic import factorized_joint_logprob
+from pymc.logprob.basic import conditional_logp
 from pymc.logprob.rewriting import (
     PreserveRVMappings,
     construct_ir_fgraph,
@@ -310,7 +310,7 @@ def logprob_ScanRV(op, values, *inputs, name=None, **kwargs):
 
     def create_inner_out_logp(value_map: Dict[TensorVariable, TensorVariable]) -> TensorVariable:
         """Create a log-likelihood inner-output for a `Scan`."""
-        logp_parts = factorized_joint_logprob(value_map, warn_rvs=False)
+        logp_parts = conditional_logp(value_map, warn_rvs=False)
         return logp_parts.values()
 
     logp_scan_args = convert_outer_out_to_in(

pymc/model.py

@@ -64,7 +64,7 @@
     ShapeWarning,
 )
 from pymc.initial_point import make_initial_point_fn
-from pymc.logprob.basic import joint_logp
+from pymc.logprob.basic import transformed_conditional_logp
 from pymc.logprob.utils import ParameterValueError
 from pymc.pytensorf import (
     PointFunc,
@@ -761,7 +761,7 @@ def logp(
 
         rv_logps: List[TensorVariable] = []
         if rvs:
-            rv_logps = joint_logp(
+            rv_logps = transformed_conditional_logp(
                 rvs=rvs,
                 rvs_to_values=self.rvs_to_values,
                 rvs_to_transforms=self.rvs_to_transforms,

pymc/testing.py

@@ -35,7 +35,7 @@
 from pymc.distributions.distribution import Distribution
 from pymc.distributions.shape_utils import change_dist_size
 from pymc.initial_point import make_initial_point_fn
-from pymc.logprob.basic import icdf, joint_logp, logcdf, logp
+from pymc.logprob.basic import icdf, logcdf, logp, transformed_conditional_logp
 from pymc.logprob.utils import ParameterValueError, find_rvs_in_graph
 from pymc.pytensorf import (
     compile_pymc,
@@ -673,7 +673,7 @@ def assert_moment_is_expected(model, expected, check_finite_logp=True):
 
     if check_finite_logp:
         logp_moment = (
-            joint_logp(
+            transformed_conditional_logp(
                 (model["x"],),
                 rvs_to_values={model["x"]: pt.constant(moment)},
                 rvs_to_transforms={},

tests/distributions/test_transform.py

@@ -25,7 +25,7 @@
 import pymc as pm
 import pymc.distributions.transforms as tr
 
-from pymc.logprob.basic import joint_logp
+from pymc.logprob.basic import transformed_conditional_logp
 from pymc.pytensorf import floatX, jacobian
 from pymc.testing import (
     Circ,
@@ -308,7 +308,7 @@ def check_transform_elementwise_logp(self, model):
         assert model.logp(x, sum=False)[0].ndim == x.ndim == jacob_det.ndim
 
         v1 = (
-            joint_logp(
+            transformed_conditional_logp(
                 (x,),
                 rvs_to_values={x: x_val_transf},
                 rvs_to_transforms={x: transform},
@@ -318,7 +318,7 @@ def check_transform_elementwise_logp(self, model):
             .eval({x_val_transf: test_array_transf})
         )
         v2 = (
-            joint_logp(
+            transformed_conditional_logp(
                 (x,),
                 rvs_to_values={x: x_val_untransf},
                 rvs_to_transforms={},
@@ -356,7 +356,7 @@ def check_vectortransform_elementwise_logp(self, model):
             assert model.logp(x, sum=False)[0].ndim == (x.ndim - 1) == jacob_det.ndim
 
         a = (
-            joint_logp(
+            transformed_conditional_logp(
                 (x,),
                 rvs_to_values={x: x_val_transf},
                 rvs_to_transforms={x: transform},
@@ -366,7 +366,7 @@ def check_vectortransform_elementwise_logp(self, model):
             .eval({x_val_transf: test_array_transf})
         )
         b = (
-            joint_logp(
+            transformed_conditional_logp(
                 (x,),
                 rvs_to_values={x: x_val_untransf},
                 rvs_to_transforms={},

tests/logprob/test_basic.py

@@ -55,7 +55,13 @@
 
 import pymc as pm
 
-from pymc.logprob.basic import factorized_joint_logprob, icdf, joint_logp, logcdf, logp
+from pymc.logprob.basic import (
+    conditional_logp,
+    icdf,
+    logcdf,
+    logp,
+    transformed_conditional_logp,
+)
 from pymc.logprob.transforms import LogTransform
 from pymc.logprob.utils import rvs_to_value_vars, walk_model
 from pymc.pytensorf import replace_rvs_by_values
@@ -68,7 +74,7 @@ def test_factorized_joint_logprob_basic():
     a.name = "a"
     a_value_var = a.clone()
 
-    a_logp = factorized_joint_logprob({a: a_value_var})
+    a_logp = conditional_logp({a: a_value_var})
     a_logp_comb = tuple(a_logp.values())[0]
     a_logp_exp = logp(a, a_value_var)
 
@@ -81,7 +87,7 @@ def test_factorized_joint_logprob_basic():
     sigma_value_var = sigma.clone()
     y_value_var = Y.clone()
 
-    total_ll = factorized_joint_logprob({Y: y_value_var, sigma: sigma_value_var})
+    total_ll = conditional_logp({Y: y_value_var, sigma: sigma_value_var})
     total_ll_combined = pt.add(*total_ll.values())
 
     # We need to replace the reference to `sigma` in `Y` with its value
@@ -106,7 +112,7 @@ def test_factorized_joint_logprob_basic():
     b_value_var = b.clone()
     c_value_var = c.clone()
 
-    b_logp = factorized_joint_logprob({a: a_value_var, b: b_value_var, c: c_value_var})
+    b_logp = conditional_logp({a: a_value_var, b: b_value_var, c: c_value_var})
     b_logp_combined = pt.sum([pt.sum(factor) for factor in b_logp.values()])
 
     # There shouldn't be any `RandomVariable`s in the resulting graph
@@ -125,7 +131,7 @@ def test_factorized_joint_logprob_multi_obs():
     a_val = a.clone()
     b_val = b.clone()
 
-    logp_res = factorized_joint_logprob({a: a_val, b: b_val})
+    logp_res = conditional_logp({a: a_val, b: b_val})
     logp_res_combined = pt.add(*logp_res.values())
     logp_exp = logp(a, a_val) + logp(b, b_val)
 
@@ -137,8 +143,8 @@ def test_factorized_joint_logprob_multi_obs():
     x_val = x.clone()
     y_val = y.clone()
 
-    logp_res = factorized_joint_logprob({x: x_val, y: y_val})
-    exp_logp = factorized_joint_logprob({x: x_val, y: y_val})
+    logp_res = conditional_logp({x: x_val, y: y_val})
+    exp_logp = conditional_logp({x: x_val, y: y_val})
     logp_res_comb = pt.sum([pt.sum(factor) for factor in logp_res.values()])
     exp_logp_comb = pt.sum([pt.sum(factor) for factor in exp_logp.values()])
 
@@ -155,7 +161,7 @@ def test_factorized_joint_logprob_diff_dims():
     y_vv = y.clone()
     y_vv.name = "y"
 
-    logp = factorized_joint_logprob({x: x_vv, y: y_vv})
+    logp = conditional_logp({x: x_vv, y: y_vv})
     logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
 
     M_val = np.random.normal(size=(10, 3))
@@ -181,7 +187,7 @@ def test_incsubtensor_original_values_output_dict():
     rv = pt.set_subtensor(base_rv[0], 5)
     vv = rv.clone()
 
-    logp_dict = factorized_joint_logprob({rv: vv})
+    logp_dict = conditional_logp({rv: vv})
     assert vv in logp_dict
 
 
@@ -194,14 +200,14 @@ def test_persist_inputs():
     beta_vv = beta_rv.type()
     y_vv = Y_rv.clone()
 
-    logp = factorized_joint_logprob({beta_rv: beta_vv, Y_rv: y_vv})
+    logp = conditional_logp({beta_rv: beta_vv, Y_rv: y_vv})
     logp_combined = pt.sum([pt.sum(factor) for factor in logp.values()])
 
     assert x in ancestors([logp_combined])
 
     # Make sure we don't clone value variables when they're graphs.
     y_vv_2 = y_vv * 2
-    logp_2 = factorized_joint_logprob({beta_rv: beta_vv, Y_rv: y_vv_2})
+    logp_2 = conditional_logp({beta_rv: beta_vv, Y_rv: y_vv_2})
     logp_2_combined = pt.sum([pt.sum(factor) for factor in logp_2.values()])
 
     assert y_vv in ancestors([logp_2_combined])
@@ -210,7 +216,7 @@ def test_persist_inputs():
     # Even when they are random
     y_vv = pt.random.normal(name="y_vv2")
     y_vv_2 = y_vv * 2
-    logp_2 = factorized_joint_logprob({beta_rv: beta_vv, Y_rv: y_vv_2})
+    logp_2 = conditional_logp({beta_rv: beta_vv, Y_rv: y_vv_2})
     logp_2_combined = pt.sum([pt.sum(factor) for factor in logp_2.values()])
 
     assert y_vv in ancestors([logp_2_combined])
@@ -224,11 +230,11 @@ def test_warn_random_found_factorized_joint_logprob():
     y_vv = y_rv.clone()
 
     with pytest.warns(UserWarning, match="Random variables detected in the logp graph: {x}"):
-        factorized_joint_logprob({y_rv: y_vv})
+        conditional_logp({y_rv: y_vv})
 
     with warnings.catch_warnings():
         warnings.simplefilter("error")
-        factorized_joint_logprob({y_rv: y_vv}, warn_missing_rvs=False)
+        conditional_logp({y_rv: y_vv}, warn_missing_rvs=False)
 
 
 def test_multiple_rvs_to_same_value_raises():
@@ -237,9 +243,9 @@ def test_multiple_rvs_to_same_value_raises():
     x = x_rv1.type()
     x.name = "x"
 
-    msg = "More than one logprob factor was assigned to the value var x"
+    msg = "More than one logprob term was assigned to the value var x"
     with pytest.raises(ValueError, match=msg):
-        factorized_joint_logprob({x_rv1: x, x_rv2: x})
+        conditional_logp({x_rv1: x, x_rv2: x})
 
 
 def test_joint_logp_basic():
@@ -259,7 +265,7 @@ def test_joint_logp_basic():
 
     c_value_var = m.rvs_to_values[c]
 
-    (b_logp,) = joint_logp(
+    (b_logp,) = transformed_conditional_logp(
         (b,),
         rvs_to_values=m.rvs_to_values,
         rvs_to_transforms=m.rvs_to_transforms,
@@ -304,7 +310,7 @@ def test_joint_logp_incsubtensor(indices, size):
     a_idx_value_var = a_idx.type()
     a_idx_value_var.name = "a_idx_value"
 
-    a_idx_logp = joint_logp(
+    a_idx_logp = transformed_conditional_logp(
         (a_idx,),
         rvs_to_values={a_idx: a_value_var},
         rvs_to_transforms={},