Implement logp for add and mul ops involving one unregistered random variable

Author: ricardoV94

pymc3/distributions/logp.py

@@ -25,6 +25,8 @@
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.op import Op, compute_test_value
 from aesara.graph.type import CType
+from aesara.scalar.basic import Add, Mul
+from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.op import RandomVariable
 from aesara.tensor.random.opt import local_subtensor_rv_lift
 from aesara.tensor.subtensor import (
@@ -37,7 +39,12 @@
 )
 from aesara.tensor.var import TensorVariable
 
-from pymc3.aesaraf import extract_rv_and_value_vars, floatX, rvs_to_value_vars
+from pymc3.aesaraf import (
+    extract_rv_and_value_vars,
+    floatX,
+    rvs_to_value_vars,
+    walk_model,
+)
 
 
 @singledispatch
@@ -260,6 +267,130 @@ def _logp(
     """
     value_var = rvs_to_values.get(var, var)
     return at.zeros_like(value_var)
+    # raise NotImplementedError(f"Logp cannot be computed for op {op}")
+
+
+@_logp.register(Elemwise)
+def logp_elemwise(op, *args, **kwargs):
+    if hasattr(op, "scalar_op"):
+        return _logp(op.scalar_op, *args, **kwargs)
+    raise NotImplementedError
+
+
+# TODO: Implement DimShuffle logp?
+# @_logp.register(DimShuffle)
+# def logp_dimshuffle(op, var, *args, **kwargs):
+#     if var.owner and len(var.owner.inputs) == 1:
+#         inp = var.owner.inputs[0]
+#         if inp.owner and hasattr(inp.owner, 'op'):
+#             return _logp(inp.owner.op, inp, *args, **kwargs)
+#     raise NotImplementedError
+
+
+def find_rv_branch(inputs):
+    """
+    Helper function to find which input branch(es) contain unregistered random variables
+    """
+    rv_branch = []
+    no_rv_branch = []
+
+    for inp in inputs:
+        res_ancestors = list(walk_model((inp,), walk_past_rvs=True))
+        # unregistered variables do not contain a value_var tag
+        res_unregistered_ancestors = [
+            v
+            for v in res_ancestors
+            if v.owner
+            and isinstance(v.owner.op, RandomVariable)
+            and not getattr(v.tag, "value_var", False)
+        ]
+        if res_unregistered_ancestors:
+            rv_branch.append(inp)
+        else:
+            no_rv_branch.append(inp)
+
+    return rv_branch, no_rv_branch
+
+
+@_logp.register(Add)
+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)
+
+    if len(rv) != 1:
+        raise NotImplementedError(
+            f"Logp of addition requires one branch with an unregistered RandomVariable but got {len(rv)}"
+        )
+
+    rv = rv[0]
+    rv_value = rvs_to_values.get(rv, getattr(rv.tag, "value_var", rv))
+    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
+
+    logp_rv = logpt(rv, new_rvs_to_values, **kwargs)
+    fgraph = FunctionGraph(
+        [i for i in graph_inputs((logp_rv,)) if not isinstance(i, Constant)],
+        [logp_rv],
+        clone=False,
+    )
+
+    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}"
+
+    return logp_rv
+
+
+@_logp.register(Mul)
+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)
+
+    if len(rv) != 1:
+        raise NotImplementedError(
+            f"Logp of product requires one branch with an unregistered RandomVariable but got {len(rv)}"
+        )
+
+    rv = rv[0]
+    rv_value = rvs_to_values.get(rv, getattr(rv.tag, "value_var", rv))
+    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
+
+    logp_rv = logpt(rv, new_rvs_to_values, **kwargs)
+    fgraph = FunctionGraph(
+        [i for i in graph_inputs((logp_rv,)) if not isinstance(i, Constant)],
+        [logp_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)
+
+    logp_rv = fgraph.outputs[0] - at.log(at.abs_(scale_value))
+    logp_rv.name = f"__logp_{var.name}"
+
+    return logp_rv
 
 
 def convert_indices(indices, entry):

pymc3/tests/test_logp.py

@@ -11,6 +11,8 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
+from contextlib import ExitStack as does_not_raise
+
 import aesara
 import aesara.tensor as at
 import numpy as np
@@ -31,7 +33,7 @@
 )
 
 from pymc3.aesaraf import floatX, walk_model
-from pymc3.distributions.continuous import Normal, Uniform
+from pymc3.distributions.continuous import Exponential, Normal, Uniform
 from pymc3.distributions.discrete import Bernoulli
 from pymc3.distributions.logp import logpt
 from pymc3.model import Model
@@ -69,6 +71,140 @@ def test_logpt_basic():
     assert a_value_var in res_ancestors
 
 
+def test_logpt_add():
+    """
+    Mare sure we can compute a log-likelihood for ``loc + Y`` where ``Y`` is an unregistered
+    random variable and ``loc`` is an tensor variable or a registered random variable
+    """
+    with Model() as m:
+        loc = Uniform("loc", 0, 1)
+        x = Normal.dist(0, 1) + loc
+        m.register_rv(x, "x")
+
+    loc_value_var = m.rvs_to_values[loc]
+    x_value_var = m.rvs_to_values[x]
+
+    x_logp = logpt(x, m.rvs_to_values[x])
+
+    res_ancestors = list(walk_model((x_logp,), walk_past_rvs=True))
+    res_rv_ancestors = [
+        v for v in res_ancestors if v.owner and isinstance(v.owner.op, RandomVariable)
+    ]
+
+    # There shouldn't be any `RandomVariable`s in the resulting graph
+    assert len(res_rv_ancestors) == 0
+    assert loc_value_var in res_ancestors
+    assert x_value_var in res_ancestors
+
+    # 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)
+
+
+def test_logpt_mul():
+    """
+    Mare sure we can compute a log-likelihood for ``scale * Y`` where ``Y`` is an unregistered
+    random variable and ``scale`` is an tensor variable or a registered random variable
+    """
+    with Model() as m:
+        scale = Uniform("scale", 0, 1)
+        x = Exponential.dist(1) * scale
+        m.register_rv(x, "x")
+
+    scale_value_var = m.rvs_to_values[scale]
+    x_value_var = m.rvs_to_values[x]
+
+    x_logp = logpt(x, m.rvs_to_values[x])
+
+    res_ancestors = list(walk_model((x_logp,), walk_past_rvs=True))
+    res_rv_ancestors = [
+        v for v in res_ancestors if v.owner and isinstance(v.owner.op, RandomVariable)
+    ]
+
+    # There shouldn't be any `RandomVariable`s in the resulting graph
+    assert len(res_rv_ancestors) == 0
+    assert scale_value_var in res_ancestors
+    assert x_value_var in res_ancestors
+
+    # 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
+
+
+def test_logpt_mul_add():
+    """
+    Mare sure we can compute a log-likelihood for ``loc + scale * Y`` where ``Y`` is an unregistered
+    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)
+        x = loc + scale * Normal.dist(0, 1)
+        m.register_rv(x, "x")
+
+    loc_value_var = m.rvs_to_values[loc]
+    scale_value_var = m.rvs_to_values[scale]
+    x_value_var = m.rvs_to_values[x]
+
+    x_logp = logpt(x, m.rvs_to_values[x])
+
+    res_ancestors = list(walk_model((x_logp,), walk_past_rvs=True))
+    res_rv_ancestors = [
+        v for v in res_ancestors if v.owner and isinstance(v.owner.op, RandomVariable)
+    ]
+
+    # There shouldn't be any `RandomVariable`s in the resulting graph
+    assert len(res_rv_ancestors) == 0
+    assert loc_value_var in res_ancestors
+    assert scale_value_var in res_ancestors
+    assert x_value_var in res_ancestors
+
+    # 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)
+
+
+def test_logpt_not_implemented():
+    """Test that logpt for add and mul fail if inputs are 0 or 2 unregistered rvs"""
+
+    with Model() as m:
+        variable1 = at.as_tensor_variable(1, "variable1")
+        variable2 = at.scalar("variable2")
+        unregistered1 = Normal.dist(0, 1)
+        unregistered2 = Normal.dist(0, 1)
+        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_pass1 = variable1 + unregistered2
+        x_pass2 = unregistered1 + variable2
+        x_pass3 = registered1 + unregistered1
+
+        m.register_rv(x_fail1, "x_fail1")
+        m.register_rv(x_fail2, "x_fail2")
+        m.register_rv(x_fail3, "x_fail3")
+        m.register_rv(x_fail4, "x_fail4")
+        m.register_rv(x_pass1, "x_pass1")
+        m.register_rv(x_pass2, "x_pass2")
+        m.register_rv(x_pass3, "x_pass3")
+
+    for rv, value_var in m.rvs_to_values.items():
+        if "fail" in rv.name:
+            with pytest.raises(NotImplementedError):
+                logpt(rv, value_var)
+        else:
+            with does_not_raise():
+                logpt(rv, value_var)
+
+
 @pytest.mark.parametrize(
     "indices, size",
     [