Create value variables when building the logprob graphs

Author: rlouf

aeppl/joint_logprob.py

@@ -1,6 +1,6 @@
 import warnings
 from collections import deque
-from typing import Dict, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import aesara.tensor as at
 from aesara import config
@@ -16,26 +16,27 @@
 
 
 def conditional_logprob(
-    rv_values: Dict[TensorVariable, TensorVariable],
+    *random_variables: TensorVariable,
+    realized: 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 conditional log-probabilities.
+) -> Tuple[Dict[TensorVariable, TensorVariable], List[TensorVariable]]:
+    r"""Create a map between random variables and their conditional log-probabilities.
 
-    The `rvs` list implicitly defines a joint probability, that factorizes
-    according to the graphical model represented by the Aesara model the
-    `RandomVariable`s belong to.
+    The list of measurable variables implicitly defines a joint probability that
+    factorizes according to the graphical model implemented by the Aesara model
+    these variables belong to.
 
     For example, consider the following
 
     .. code-block:: python
 
         import aesara.tensor as at
 
-        sigma2_rv = at.random.invgamma(0.5, 0.5)
-        Y_rv = at.random.normal(0, at.sqrt(sigma2_rv))
+        sigma2_rv = at.random.invgamma(0.5, 0.5, name="sigma2")
+        Y_rv = at.random.normal(0, at.sqrt(sigma2_rv), name="Y")
 
     This graph for ``Y_rv`` is equivalent to the following hierarchical model:
 
@@ -51,9 +52,9 @@ def conditional_logprob(
     :math:`Y`'s respective conditional log-probabilities, :math:`\log P(\sigma^2 = s)`
     and :math:`\log p(Y = y | \sigma^2 = s)`.
 
-    `conditional_logprob` generates the value variables that correspond to the
-    measurable variables for which it produces a conditional log-probability
-    graph and returns them along with the graphs:
+    To build the log-probability graphs, `conditional_logprob` must generate
+    value variables associated with each input variable. They are returned along
+    with the graphs:
 
     .. code-block:: python
 
@@ -62,16 +63,18 @@ def conditional_logprob(
         sigma2_rv = at.random.invgamma(0.5, 0.5)
         Y_rv = at.random.normal(0, at.sqrt(sigma2_rv))
 
-        logprobs = conditional_logprob(Y_rv, sigma2_rv)
+        logprobs, value_variables = conditional_logprob(Y_rv, sigma2_rv)
         # print(logprobs.keys())
-        # [sigma2_vv, Y_vv]
+        # [Y, sigma2]
+        # print(value_variables)
+        # [Y_vv, sigma2_vv]
 
 
     Parameters
     ==========
-    rv_values
-        A ``dict`` that maps measurable variables (e.g. `RandomVariable`s) to
-        symbolic `Variable`\s that represent their values.
+    random_variables
+        A ``list`` of  random variables for which we need to return a
+        conditional log-probability graph.
     warn_missing_rvs
         When ``True``, issue a warning when a `RandomVariable` is found in
         the graph and doesn't have a corresponding value variable specified in
@@ -84,11 +87,27 @@ def conditional_logprob(
 
     Returns
     =======
-    A ``dict`` that maps each value variable to the log-probability factor derived
-    from the respective `RandomVariable`.
+    A ``dict`` that maps each random variable to the derived log-probability
+    factor, and a list of the created valued variables in the same order as the
+    order in which their corresponding random variables were passed as
+    arguments.
 
     """
+
+    # Create value variables by cloning the input measurable variables
+    original_rv_values = {}
+    for rv in random_variables:
+        vv = rv.clone()
+        if rv.name:
+            vv.name = f"{rv.name}_vv"
+        original_rv_values[rv] = vv
+
+    # Value variables are not cloned when constructing the conditional log-proprobability
+    # graphs. We can thus use them to recover the original random variables to index the
+    # maps to the logprob graphs and value variables before returning them.
+    rv_values = {**original_rv_values, **realized}
     vv_to_original_rvs = {vv: rv for rv, vv in rv_values.items()}
+
     fgraph, rv_values, _ = construct_ir_fgraph(rv_values, ir_rewriter=ir_rewriter)
 
     if extra_rewrites is not None:
@@ -194,30 +213,45 @@ def conditional_logprob(
             f"The logprob terms of the following random variables could not be derived: {missing_value_terms}"
         )
 
-    return logprob_vars
+    return logprob_vars, list(original_rv_values.values())
 
 
 def joint_logprob(
-    rv_values: Dict[TensorVariable, TensorVariable], *args, **kwargs
-) -> Optional[TensorVariable]:
+    *random_variables: List[TensorVariable],
+    realized: Dict[TensorVariable, TensorVariable] = {},
+    **kwargs,
+) -> Optional[Tuple[TensorVariable, List[TensorVariable]]]:
     """Create a graph representing the joint log-probability/measure of a graph.
 
     This function calls `factorized_joint_logprob` and returns the combined
     log-probability factors as a single graph.
 
     Parameters
-    ----------
-    sum: bool
-        If ``True`` each factor is collapsed to a scalar via ``sum`` before
-        being joined with the remaining factors. This may be necessary to
-        avoid incorrect broadcasting among independent factors.
+    ==========
+    random_variables
+        A ``list`` of  random variables for which we need to return a
+        conditional log-probability graph.
+    realized
+        A ``dict`` that maps  random variables to their realized value.
+
+    Returns
+    =======
+    A ``TensorVariable`` that represents the joint log-probability of the graph
+    implicitly defined by the random variables passed as arguments, and a list
+    of the created valued variables in the same order as the order in which
+    their corresponding random variables were passed as arguments.
 
     """
-    logprob = conditional_logprob(rv_values, *args, **kwargs)
+    logprob, value_variables = conditional_logprob(
+        *random_variables, realized=realized, **kwargs
+    )
     if not logprob:
         return None
     elif len(logprob) == 1:
-        logprob = tuple(logprob.values())[0]
-        return at.sum(logprob)
+        cond_logprob = tuple(logprob.values())[0]
+        return at.sum(cond_logprob), value_variables
     else:
-        return at.sum([at.sum(factor) for factor in logprob.values()])
+        joint_logprob: TensorVariable = at.sum(
+            [at.sum(factor) for factor in logprob.values()]
+        )
+        return joint_logprob, value_variables

aeppl/scan.py

@@ -268,7 +268,7 @@ def create_inner_out_logp(
         value_map: Dict[TensorVariable, TensorVariable]
     ) -> TensorVariable:
         """Create a log-likelihood inner-output for a `Scan`."""
-        logp_parts = conditional_logprob(value_map, warn_missing_rvs=False)
+        logp_parts, _ = conditional_logprob(realized=value_map, warn_missing_rvs=False)
         return logp_parts.values()
 
     logp_scan_args = convert_outer_out_to_in(

tests/test_censoring.py

@@ -15,13 +15,10 @@ def test_continuous_rv_clip():
     x_rv = at.random.normal(0.5, 1)
     cens_x_rv = at.clip(x_rv, -2, 2)
 
-    cens_x_vv = cens_x_rv.clone()
-    cens_x_vv.tag.test_value = 0
-
-    logp = joint_logprob({cens_x_rv: cens_x_vv})
+    logp, vv = joint_logprob(cens_x_rv)
     assert_no_rvs(logp)
 
-    logp_fn = aesara.function([cens_x_vv], logp)
+    logp_fn = aesara.function(vv, logp)
     ref_scipy = st.norm(0.5, 1)
 
     assert logp_fn(-3) == -np.inf
@@ -36,12 +33,10 @@ def test_discrete_rv_clip():
     x_rv = at.random.poisson(2)
     cens_x_rv = at.clip(x_rv, 1, 4)
 
-    cens_x_vv = cens_x_rv.clone()
-
-    logp = joint_logprob({cens_x_rv: cens_x_vv})
+    logp, vv = joint_logprob(cens_x_rv)
     assert_no_rvs(logp)
 
-    logp_fn = aesara.function([cens_x_vv], logp)
+    logp_fn = aesara.function(vv, logp)
     ref_scipy = st.poisson(2)
 
     assert logp_fn(0) == -np.inf
@@ -57,11 +52,8 @@ def test_one_sided_clip():
     lb_cens_x_rv = at.clip(x_rv, -1, x_rv)
     ub_cens_x_rv = at.clip(x_rv, x_rv, 1)
 
-    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, (lb_cens_x_vv,) = joint_logprob(lb_cens_x_rv)
+    ub_logp, (ub_cens_x_vv,) = joint_logprob(ub_cens_x_rv)
     assert_no_rvs(lb_logp)
     assert_no_rvs(ub_logp)
 
@@ -78,9 +70,8 @@ def test_useless_clip():
     x_rv = at.random.normal(0.5, 1, size=3)
     cens_x_rv = at.clip(x_rv, x_rv, x_rv)
 
-    cens_x_vv = cens_x_rv.clone()
-
-    logp = conditional_logprob({cens_x_rv: cens_x_vv})[cens_x_rv]
+    logps, (cens_x_vv,) = conditional_logprob(cens_x_rv)
+    logp = logps[cens_x_rv]
     assert_no_rvs(logp)
 
     logp_fn = aesara.function([cens_x_vv], logp)
@@ -94,9 +85,7 @@ def test_random_clip():
     x_rv = at.random.normal(0, 2)
     cens_x_rv = at.clip(x_rv, lb_rv, [1, 1])
 
-    lb_vv = lb_rv.clone()
-    cens_x_vv = cens_x_rv.clone()
-    logps = conditional_logprob({cens_x_rv: cens_x_vv, lb_rv: lb_vv})
+    logps, (cens_x_vv, lb_vv) = conditional_logprob(cens_x_rv, lb_rv)
     logp = at.add(*logps.values())
     assert_no_rvs(logp)
 
@@ -111,10 +100,7 @@ def test_broadcasted_clip_constant():
     x_rv = at.random.normal(0, 2)
     cens_x_rv = at.clip(x_rv, lb_rv, [1, 1])
 
-    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})
+    logp, _ = joint_logprob(cens_x_rv, lb_rv)
     assert_no_rvs(logp)
 
 
@@ -123,10 +109,7 @@ def test_broadcasted_clip_random():
     x_rv = at.random.normal(0, 2, size=2)
     cens_x_rv = at.clip(x_rv, lb_rv, 1)
 
-    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})
+    logp, _ = joint_logprob(cens_x_rv, lb_rv)
     assert_no_rvs(logp)
 
 
@@ -136,10 +119,8 @@ def test_fail_base_and_clip_have_values():
     cens_x_rv = at.clip(x_rv, x_rv, 1)
     cens_x_rv.name = "cens_x"
 
-    x_vv = x_rv.clone()
-    cens_x_vv = cens_x_rv.clone()
     with pytest.raises(RuntimeError, match="could not be derived: {cens_x}"):
-        conditional_logprob({cens_x_rv: cens_x_vv, x_rv: x_vv})
+        conditional_logprob(cens_x_rv, x_rv)
 
 
 def test_fail_multiple_clip_single_base():
@@ -150,20 +131,16 @@ def test_fail_multiple_clip_single_base():
     cens_rv2 = at.clip(base_rv, -1, 1)
     cens_rv2.name = "cens2"
 
-    cens_vv1 = cens_rv1.clone()
-    cens_vv2 = cens_rv2.clone()
     with pytest.raises(RuntimeError, match="could not be derived: {cens2}"):
-        conditional_logprob({cens_rv1: cens_vv1, cens_rv2: cens_vv2})
+        conditional_logprob(cens_rv1, cens_rv2)
 
 
 def test_deterministic_clipping():
     x_rv = at.random.normal(0, 1)
     clip = at.clip(x_rv, 0, 0)
     y_rv = at.random.normal(clip, 1)
 
-    x_vv = x_rv.clone()
-    y_vv = y_rv.clone()
-    logp = joint_logprob({x_rv: x_vv, y_rv: y_vv})
+    logp, (x_vv, y_vv) = joint_logprob(x_rv, y_rv)
     assert_no_rvs(logp)
 
     logp_fn = aesara.function([x_vv, y_vv], logp)
@@ -180,7 +157,7 @@ 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, _ = joint_logprob(realized={cens_x_rv: cens_x_vv}, extra_rewrites=transform)
 
     cens_x_vv_testval = -1
     obs_logp = logp.eval({cens_x_vv: cens_x_vv_testval})
@@ -201,8 +178,8 @@ def test_rounding(rounding_op):
     xr = rounding_op(x)
     xr.name = "xr"
 
-    xr_vv = xr.clone()
-    logp = conditional_logprob({xr: xr_vv})[xr]
+    logp, (xr_vv,) = conditional_logprob(xr)
+    logp = logp[xr]
     assert logp is not None
 
     x_sp = st.norm(loc, scale)