Make mul_logp and add_logp work with transforms

Author: ricardoV94

pymc3/distributions/logp.py

@@ -318,37 +318,38 @@ def add_logp(op, var, rvs_to_values, *add_inputs, **kwargs):
     if len(add_inputs) != 2:
         raise ValueError(f"Expected 2 inputs but got: {len(add_inputs)}")
 
-    rv, loc = find_rv_branch(add_inputs)
+    base_rv, loc = find_rv_branch(add_inputs)
 
-    if len(rv) != 1:
+    if len(base_rv) != 1:
         raise NotImplementedError(
-            f"Logp of addition requires one branch with an unregistered RandomVariable but got {len(rv)}"
+            f"Logp of addition requires one branch with an unregistered RandomVariable but got {len(base_rv)}"
         )
 
-    rv = rv[0]
-    rv_value = rvs_to_values.get(rv, getattr(rv.tag, "value_var", rv))
+    var_value = rvs_to_values.get(var, var)
     loc = loc[0]
-    loc_value = rvs_to_values.get(loc, getattr(loc.tag, "value_var", loc))
-
-    new_rvs_to_values = rvs_to_values.copy()
-    new_rvs_to_values[rv] = rv_value
+    base_rv = base_rv[0]
+    base_value = base_rv.type()
 
-    logp_rv = logpt(rv, new_rvs_to_values, **kwargs)
+    logp_base_rv = logpt(base_rv, {base_rv: base_value}, **kwargs)
     fgraph = FunctionGraph(
-        [i for i in graph_inputs((logp_rv,)) if not isinstance(i, Constant)],
-        [logp_rv],
+        [i for i in graph_inputs((logp_base_rv,)) if not isinstance(i, Constant)],
+        [logp_base_rv],
         clone=False,
     )
+    fgraph.replace(base_value, var_value - loc, import_missing=True)
+    logp_add_rv = fgraph.outputs[0]
+
+    # Replace rvs in graph
+    # TODO: This shouldn't be here
+    (logp_add_rv,), _ = rvs_to_value_vars(
+        (logp_add_rv,),
+        apply_transforms=True,  # Change this
+        initial_replacements=None,
+    )
 
-    var_value = rvs_to_values.get(var, var)
-
-    fgraph.add_input(loc_value)
-    fgraph.add_input(var_value)
-    fgraph.replace(rv_value, var_value - loc_value)
-
-    logp_rv.name = f"__logp_{var.name}"
+    logp_add_rv.name = f"__logp_{var.name}"
 
-    return logp_rv
+    return logp_add_rv
 
 
 @_logp.register(Mul)
@@ -357,40 +358,41 @@ def mul_logp(op, var, rvs_to_values, *mul_inputs, **kwargs):
     if len(mul_inputs) != 2:
         raise ValueError(f"Expected 2 inputs but got: {len(mul_inputs)}")
 
-    rv, scale = find_rv_branch(mul_inputs)
+    base_rv, scale = find_rv_branch(mul_inputs)
 
-    if len(rv) != 1:
+    if len(base_rv) != 1:
         raise NotImplementedError(
-            f"Logp of product requires one branch with an unregistered RandomVariable but got {len(rv)}"
+            f"Logp of product requires one branch with an unregistered RandomVariable but got {len(base_rv)}"
         )
 
-    rv = rv[0]
-    rv_value = rvs_to_values.get(rv, getattr(rv.tag, "value_var", rv))
+    var_value = rvs_to_values.get(var, var)
     scale = scale[0]
-    scale_value = rvs_to_values.get(scale, getattr(scale.tag, "value_var", scale))
-
-    new_rvs_to_values = rvs_to_values.copy()
-    new_rvs_to_values[rv] = rv_value
+    base_rv = base_rv[0]
+    base_value = base_rv.type()
 
-    logp_rv = logpt(rv, new_rvs_to_values, **kwargs)
+    logp_base_rv = logpt(base_rv, {base_rv: base_value}, **kwargs)
     fgraph = FunctionGraph(
-        [i for i in graph_inputs((logp_rv,)) if not isinstance(i, Constant)],
-        [logp_rv],
+        [i for i in graph_inputs((logp_base_rv,)) if not isinstance(i, Constant)],
+        [logp_base_rv],
         clone=False,
     )
 
-    var_value = rvs_to_values.get(var, var)
-
-    fgraph.add_input(scale_value)
-    fgraph.add_input(var_value)
     # TODO: This is not correct for discrete variables
     # TODO: Undefined behavior for scale = 0
-    fgraph.replace(rv_value, var_value / scale_value)
+    fgraph.replace(base_value, var_value / scale, import_missing=True)
+    logp_mul_rv = fgraph.outputs[0] - at.log(at.abs_(scale))
+
+    # Replace rvs in graph
+    # TODO: This shouldn't be here
+    (logp_mul_rv,), _ = rvs_to_value_vars(
+        (logp_mul_rv,),
+        apply_transforms=True,  # Change this
+        initial_replacements=None,
+    )
 
-    logp_rv = fgraph.outputs[0] - at.log(at.abs_(scale_value))
-    logp_rv.name = f"__logp_{var.name}"
+    logp_mul_rv.name = f"__logp_{var.name}"
 
-    return logp_rv
+    return logp_mul_rv
 
 
 def convert_indices(indices, entry):

pymc3/tests/test_logp.py

@@ -77,7 +77,7 @@ def test_logpt_add():
     random variable and ``loc`` is an tensor variable or a registered random variable
     """
     with Model() as m:
-        loc = Uniform("loc", 0, 1)
+        loc = Exponential("loc", 10)
         x = Normal.dist(0, 1) + loc
         m.register_rv(x, "x")
 
@@ -98,8 +98,8 @@ def test_logpt_add():
 
     # Test logp is correct
     f_logp = aesara.function([x_value_var, loc_value_var], x_logp)
-    np.testing.assert_almost_equal(f_logp(50, 50), sp.norm(50, 1).logpdf(50))
-    np.testing.assert_almost_equal(f_logp(50, 0), sp.norm(0, 1).logpdf(50), decimal=5)
+    np.testing.assert_almost_equal(f_logp(50, np.log(50)), sp.norm(50, 1).logpdf(50))
+    np.testing.assert_almost_equal(f_logp(50, np.log(10)), sp.norm(10, 1).logpdf(50), decimal=5)
 
 
 def test_logpt_mul():
@@ -108,7 +108,7 @@ def test_logpt_mul():
     random variable and ``scale`` is an tensor variable or a registered random variable
     """
     with Model() as m:
-        scale = Uniform("scale", 0, 1)
+        scale = Exponential("scale", 10)
         x = Exponential.dist(1) * scale
         m.register_rv(x, "x")
 
@@ -129,9 +129,8 @@ def test_logpt_mul():
 
     # Test logp is correct
     f_logp = aesara.function([x_value_var, scale_value_var], x_logp)
-    np.testing.assert_almost_equal(f_logp(0, 5), sp.expon(scale=5).logpdf(0))
-    np.testing.assert_almost_equal(f_logp(-2, -2), sp.expon(scale=2).logpdf(2))
-    assert f_logp(2, -2) == -np.inf
+    np.testing.assert_almost_equal(f_logp(0, np.log(5)), sp.expon(scale=5).logpdf(0))
+    np.testing.assert_almost_equal(f_logp(2, np.log(2)), sp.expon(scale=2).logpdf(2))
 
 
 def test_logpt_mul_add():
@@ -140,8 +139,8 @@ def test_logpt_mul_add():
     random variable and ``loc`` and ``scale`` are tensor variables or registered random variables
     """
     with Model() as m:
-        loc = Uniform("loc", 0, 1)
-        scale = Uniform("scale", 0, 1)
+        loc = Exponential("loc", 10)
+        scale = Exponential("scale", 10)
         x = loc + scale * Normal.dist(0, 1)
         m.register_rv(x, "x")
 
@@ -164,11 +163,14 @@ def test_logpt_mul_add():
 
     # Test logp is correct
     f_logp = aesara.function([x_value_var, loc_value_var, scale_value_var], x_logp)
-    np.testing.assert_almost_equal(f_logp(-1, 0, 2), sp.norm(0, 2).logpdf(-1))
-    np.testing.assert_almost_equal(f_logp(95, 100, 15), sp.norm(100, 15).logpdf(95), decimal=6)
+    np.testing.assert_almost_equal(f_logp(-1, np.log(0), np.log(2)), sp.norm(0, 2).logpdf(-1))
+    np.testing.assert_almost_equal(
+        f_logp(95, np.log(100), np.log(15)), sp.norm(100, 15).logpdf(95), decimal=6
+    )
 
 
-def test_logpt_not_implemented():
+@pytest.mark.parametrize("op", [at.add, at.mul])
+def test_logpt_not_implemented(op):
     """Test that logpt for add and mul fail if inputs are 0 or 2 unregistered rvs"""
 
     with Model() as m:
@@ -179,14 +181,14 @@ def test_logpt_not_implemented():
         registered1 = Normal("registered1", 0, 1)
         registered2 = Normal("registered2", 0, 1)
 
-        x_fail1 = variable1 + variable2
-        x_fail2 = unregistered1 + unregistered2
-        x_fail3 = registered1 + variable1
-        x_fail4 = registered1 + registered2
+        x_fail1 = op(variable1, variable2)
+        x_fail2 = op(unregistered1, unregistered2)
+        x_fail3 = op(registered1, variable1)
+        x_fail4 = op(registered1, registered2)
 
-        x_pass1 = variable1 + unregistered2
-        x_pass2 = unregistered1 + variable2
-        x_pass3 = registered1 + unregistered1
+        x_pass1 = op(variable1, unregistered2)
+        x_pass2 = op(unregistered1, variable2)
+        x_pass3 = op(registered1, unregistered1)
 
         m.register_rv(x_fail1, "x_fail1")
         m.register_rv(x_fail2, "x_fail2")