Refactor logpt and raise more informative ValueErrors

ricardoV94 · ricardoV94 · commit afad6acc405c · 2021-12-12T20:12:06.000+01:00
diff --git a/pymc/distributions/logprob.py b/pymc/distributions/logprob.py
@@ -24,10 +24,8 @@
 from aeppl.logprob import logcdf as logcdf_aeppl
 from aeppl.logprob import logprob as logp_aeppl
 from aeppl.transforms import TransformValuesOpt
-from aesara import config
 from aesara.graph.basic import graph_inputs, io_toposort
-from aesara.graph.op import Op, compute_test_value
-from aesara.tensor.random.op import RandomVariable
+from aesara.graph.op import Op
 from aesara.tensor.subtensor import (
     AdvancedIncSubtensor,
     AdvancedIncSubtensor1,
@@ -164,100 +162,86 @@ def logpt(
     # joint_logprob directly.
 
     # If var is not a list make it one.
-    if not isinstance(var, list):
+    if not isinstance(var, (list, tuple)):
         var = [var]
 
-    # If logpt isn't provided values and the variable (provided in var)
-    # is an RV, it is assumed that the tagged value var or observation is
-    # the value variable for that particular RV.
+    # If logpt isn't provided values it is assumed that the tagged value var or
+    # observation is the value variable for that particular RV.
     if rv_values is None:
         rv_values = {}
-        for _var in var:
-            if isinstance(_var.owner.op, RandomVariable):
-                rv_value_var = getattr(
-                    _var.tag, "observations", getattr(_var.tag, "value_var", _var)
-                )
-                rv_values = {_var: rv_value_var}
+        for rv in var:
+            value_var = getattr(rv.tag, "observations", getattr(rv.tag, "value_var", None))
+            if value_var is None:
+                raise ValueError(f"No value variable found for var {rv}")
+            rv_values[rv] = value_var
+    # Else we assume we were given a single rv and respective value
     elif not isinstance(rv_values, Mapping):
-        # Else if we're given a single value and a single variable we assume a mapping among them.
-        rv_values = (
-            {var[0]: at.as_tensor_variable(rv_values).astype(var[0].type)} if len(var) == 1 else {}
-        )
-
-    # Since the filtering of logp graph is based on value variables
-    # provided to this function
-    if not rv_values:
-        warnings.warn("No value variables provided the logp will be an empty graph")
+        if len(var) == 1:
+            rv_values = {var[0]: at.as_tensor_variable(rv_values).astype(var[0].type)}
+        else:
+            raise ValueError("rv_values must be a dict if more than one var is requested")
 
     if scaling:
         rv_scalings = {}
-        for _var in var:
-            rv_value_var = getattr(_var.tag, "observations", getattr(_var.tag, "value_var", _var))
-            rv_scalings[rv_value_var] = _get_scaling(
-                getattr(_var.tag, "total_size", None), rv_value_var.shape, rv_value_var.ndim
+        for rv, value_var in rv_values.items():
+            rv_scalings[value_var] = _get_scaling(
+                getattr(rv.tag, "total_size", None), value_var.shape, value_var.ndim
             )
 
     # Aeppl needs all rv-values pairs, not just that of the requested var.
     # Hence we iterate through the graph to collect them.
     tmp_rvs_to_values = rv_values.copy()
-    transform_map = {}
     for node in io_toposort(graph_inputs(var), var):
         try:
             curr_vars = [node.default_output()]
         except ValueError:
             curr_vars = node.outputs
         for curr_var in curr_vars:
-            rv_value_var = getattr(
+            if curr_var in tmp_rvs_to_values:
+                continue
+            # Check if variable has a value variable
+            value_var = getattr(
                 curr_var.tag, "observations", getattr(curr_var.tag, "value_var", None)
             )
-            if rv_value_var is None:
-                continue
-            rv_value = rv_values.get(curr_var, rv_value_var)
-            tmp_rvs_to_values[curr_var] = rv_value
-            # Along with value variables we also check for transforms if any.
-            if hasattr(rv_value_var.tag, "transform") and transformed:
-                transform_map[rv_value] = rv_value_var.tag.transform
+            if value_var is not None:
+                tmp_rvs_to_values[curr_var] = value_var
+
+    # After collecting all necessary rvs and values, we check for any value transforms
+    transform_map = {}
+    if transformed:
+        for rv, value_var in tmp_rvs_to_values.items():
+            if hasattr(value_var.tag, "transform"):
+                transform_map[value_var] = value_var.tag.transform
+            # If the provided value_variable does not have transform information, we
+            # check if the original `rv.tag.value_var` does.
+            # TODO: This logic should be replaced by an explicit dict of
+            #  `{value_var: transform}` similar to `rv_values`.
+            else:
+                original_value_var = getattr(rv.tag, "value_var", None)
+                if original_value_var is not None and hasattr(original_value_var.tag, "transform"):
+                    transform_map[value_var] = original_value_var.tag.transform
 
     transform_opt = TransformValuesOpt(transform_map)
     temp_logp_var_dict = factorized_joint_logprob(
         tmp_rvs_to_values, extra_rewrites=transform_opt, use_jacobian=jacobian, **kwargs
     )
 
     # aeppl returns the logpt for every single value term we provided to it. This includes
-    # the extra values we plugged in above so we need to filter those out.
+    # the extra values we plugged in above, so we filter those we actually wanted in the
+    # same order they were given in.
     logp_var_dict = {}
-    for value_var, _logp in temp_logp_var_dict.items():
-        if value_var in rv_values.values():
-            logp_var_dict[value_var] = _logp
+    for value_var in rv_values.values():
+        logp_var_dict[value_var] = temp_logp_var_dict[value_var]
 
-    # If it's an empty dictionary the logp is None
-    if not logp_var_dict:
-        logp_var = None
-    else:
-        # Otherwise apply appropriate scalings and at.add and/or at.sum the
-        # graphs accordingly.
-        if scaling:
-            for _value in logp_var_dict.keys():
-                if _value in rv_scalings:
-                    logp_var_dict[_value] *= rv_scalings[_value]
-
-        if len(logp_var_dict) == 1:
-            logp_var_dict = tuple(logp_var_dict.values())[0]
-            if sum:
-                logp_var = at.sum(logp_var_dict)
-            else:
-                logp_var = logp_var_dict
-        else:
-            if sum:
-                logp_var = at.sum([at.sum(factor) for factor in logp_var_dict.values()])
-            else:
-                logp_var = at.add(*logp_var_dict.values())
+    if scaling:
+        for value_var in logp_var_dict.keys():
+            if value_var in rv_scalings:
+                logp_var_dict[value_var] *= rv_scalings[value_var]
 
-        # Recompute test values for the changes introduced by the replacements
-        # above.
-        if config.compute_test_value != "off":
-            for node in io_toposort(graph_inputs((logp_var,)), (logp_var,)):
-                compute_test_value(node)
+    if sum:
+        logp_var = at.sum([at.sum(factor) for factor in logp_var_dict.values()])
+    else:
+        logp_var = at.add(*logp_var_dict.values())
 
     return logp_var
 
diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py
@@ -2521,9 +2521,11 @@ def test_continuous(self):
         assert logpt(InfBoundedNormal, 0).eval() != -np.inf
         assert logpt(InfBoundedNormal, 11).eval() != -np.inf
 
-        value = at.dscalar("x")
+        value = model.rvs_to_values[LowerNormalTransform]
         assert logpt(LowerNormalTransform, value).eval({value: -1}) != -np.inf
+        value = model.rvs_to_values[UpperNormalTransform]
         assert logpt(UpperNormalTransform, value).eval({value: 1}) != -np.inf
+        value = model.rvs_to_values[BoundedNormalTransform]
         assert logpt(BoundedNormalTransform, value).eval({value: 0}) != -np.inf
         assert logpt(BoundedNormalTransform, value).eval({value: 11}) != -np.inf
 
diff --git a/pymc/tests/test_distributions_random.py b/pymc/tests/test_distributions_random.py
@@ -45,7 +45,7 @@ def random_polyagamma(*args, **kwargs):
 from pymc.distributions.continuous import get_tau_sigma, interpolated
 from pymc.distributions.discrete import _OrderedLogistic, _OrderedProbit
 from pymc.distributions.dist_math import clipped_beta_rvs
-from pymc.distributions.logprob import logpt
+from pymc.distributions.logprob import logp
 from pymc.distributions.multivariate import _OrderedMultinomial, quaddist_matrix
 from pymc.distributions.shape_utils import to_tuple
 from pymc.tests.helpers import SeededTest, select_by_precision
@@ -1626,8 +1626,8 @@ def test_errors(self):
                 rowcov=np.eye(3),
                 colcov=np.eye(3),
             )
-            with pytest.raises(TypeError):
-                logpt(matrixnormal, aesara.tensor.ones((3, 3, 3)))
+            with pytest.raises(ValueError):
+                logp(matrixnormal, aesara.tensor.ones((3, 3, 3)))
 
         with pm.Model():
             with pytest.warns(FutureWarning):
@@ -1856,7 +1856,7 @@ def test_density_dist_without_random(self):
             pm.DensityDist(
                 "density_dist",
                 mu,
-                logp=lambda value, mu: logpt(pm.Normal.dist(mu, 1, size=100), value),
+                logp=lambda value, mu: logp(pm.Normal.dist(mu, 1, size=100), value),
                 observed=np.random.randn(100),
                 initval=0,
             )
diff --git a/pymc/tests/test_transforms.py b/pymc/tests/test_transforms.py
@@ -23,7 +23,7 @@
 import pymc as pm
 import pymc.distributions.transforms as tr
 
-from pymc.aesaraf import jacobian
+from pymc.aesaraf import floatX, jacobian
 from pymc.distributions import logpt
 from pymc.tests.checks import close_to, close_to_logical
 from pymc.tests.helpers import SeededTest
@@ -285,40 +285,46 @@ def build_model(self, distfam, params, size, transform, initval=None):
 
     def check_transform_elementwise_logp(self, model):
         x = model.free_RVs[0]
-        x0 = x.tag.value_var
-        assert x.ndim == logpt(x, sum=False).ndim
+        x_val_transf = x.tag.value_var
 
-        pt = model.initial_point
-        array = np.random.randn(*pt[x0.name].shape)
-        transform = x0.tag.transform
-        logp_notrans = logpt(x, transform.backward(array, *x.owner.inputs), transformed=False)
+        pt = model.recompute_initial_point(0)
+        test_array_transf = floatX(np.random.randn(*pt[x_val_transf.name].shape))
+        transform = x_val_transf.tag.transform
+        test_array_untransf = transform.backward(test_array_transf, *x.owner.inputs).eval()
 
-        jacob_det = transform.log_jac_det(aesara.shared(array), *x.owner.inputs)
-        assert logpt(x, sum=False).ndim == jacob_det.ndim
+        # Create input variable with same dimensionality as untransformed test_array
+        x_val_untransf = at.constant(test_array_untransf).type()
 
-        v1 = logpt(x, array, jacobian=False).eval()
-        v2 = logp_notrans.eval()
+        jacob_det = transform.log_jac_det(test_array_transf, *x.owner.inputs)
+        assert logpt(x, sum=False).ndim == x.ndim == jacob_det.ndim
+
+        v1 = logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf})
+        v2 = logpt(x, x_val_untransf, transformed=False).eval({x_val_untransf: test_array_untransf})
         close_to(v1, v2, tol)
 
-    def check_vectortransform_elementwise_logp(self, model, vect_opt=0):
+    def check_vectortransform_elementwise_logp(self, model):
         x = model.free_RVs[0]
-        x0 = x.tag.value_var
-        # TODO: For some reason the ndim relations
-        # dont hold up here. But final log-probablity
-        # values are what we expected.
-        # assert (x.ndim - 1) == logpt(x, sum=False).ndim
-
-        pt = model.initial_point
-        array = np.random.randn(*pt[x0.name].shape)
-        transform = x0.tag.transform
-        logp_nojac = logpt(x, transform.backward(array, *x.owner.inputs), transformed=False)
-
-        jacob_det = transform.log_jac_det(aesara.shared(array), *x.owner.inputs)
-        # assert logpt(x).ndim == jacob_det.ndim
-
+        x_val_transf = x.tag.value_var
+
+        pt = model.recompute_initial_point(0)
+        test_array_transf = floatX(np.random.randn(*pt[x_val_transf.name].shape))
+        transform = x_val_transf.tag.transform
+        test_array_untransf = transform.backward(test_array_transf, *x.owner.inputs).eval()
+
+        # Create input variable with same dimensionality as untransformed test_array
+        x_val_untransf = at.constant(test_array_untransf).type()
+
+        jacob_det = transform.log_jac_det(test_array_transf, *x.owner.inputs)
+        # Original distribution is univariate
+        if x.owner.op.ndim_supp == 0:
+            assert logpt(x, sum=False).ndim == x.ndim == (jacob_det.ndim + 1)
+        # Original distribution is multivariate
+        else:
+            assert logpt(x, sum=False).ndim == (x.ndim - 1) == jacob_det.ndim
+
+        a = logpt(x, x_val_transf, jacobian=False).eval({x_val_transf: test_array_transf})
+        b = logpt(x, x_val_untransf, transformed=False).eval({x_val_untransf: test_array_untransf})
         # Hack to get relative tolerance
-        a = logpt(x, array.astype(aesara.config.floatX), jacobian=False).eval()
-        b = logp_nojac.eval()
         close_to(a, b, np.abs(0.5 * (a + b) * tol))
 
     @pytest.mark.parametrize(
@@ -406,7 +412,7 @@ def test_vonmises(self, mu, kappa, size):
     )
     def test_dirichlet(self, a, size):
         model = self.build_model(pm.Dirichlet, {"a": a}, size=size, transform=tr.simplex)
-        self.check_vectortransform_elementwise_logp(model, vect_opt=1)
+        self.check_vectortransform_elementwise_logp(model)
 
     def test_normal_ordered(self):
         model = self.build_model(
@@ -416,7 +422,7 @@ def test_normal_ordered(self):
             initval=np.asarray([-1.0, 1.0, 4.0]),
             transform=tr.ordered,
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=0)
+        self.check_vectortransform_elementwise_logp(model)
 
     @pytest.mark.parametrize(
         "sd,size",
@@ -434,7 +440,7 @@ def test_half_normal_ordered(self, sd, size):
             initval=initval,
             transform=tr.Chain([tr.log, tr.ordered]),
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=0)
+        self.check_vectortransform_elementwise_logp(model)
 
     @pytest.mark.parametrize("lam,size", [(2.5, (2,)), (np.ones(3), (4, 3))])
     def test_exponential_ordered(self, lam, size):
@@ -446,7 +452,7 @@ def test_exponential_ordered(self, lam, size):
             initval=initval,
             transform=tr.Chain([tr.log, tr.ordered]),
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=0)
+        self.check_vectortransform_elementwise_logp(model)
 
     @pytest.mark.parametrize(
         "a,b,size",
@@ -468,7 +474,7 @@ def test_beta_ordered(self, a, b, size):
             initval=initval,
             transform=tr.Chain([tr.logodds, tr.ordered]),
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=0)
+        self.check_vectortransform_elementwise_logp(model)
 
     @pytest.mark.parametrize(
         "lower,upper,size",
@@ -491,7 +497,7 @@ def transform_params(*inputs):
             initval=initval,
             transform=tr.Chain([interval, tr.ordered]),
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=1)
+        self.check_vectortransform_elementwise_logp(model)
 
     @pytest.mark.parametrize("mu,kappa,size", [(0.0, 1.0, (2,)), (np.zeros(3), np.ones(3), (4, 3))])
     def test_vonmises_ordered(self, mu, kappa, size):
@@ -503,7 +509,7 @@ def test_vonmises_ordered(self, mu, kappa, size):
             initval=initval,
             transform=tr.Chain([tr.circular, tr.ordered]),
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=0)
+        self.check_vectortransform_elementwise_logp(model)
 
     @pytest.mark.parametrize(
         "lower,upper,size,transform",
@@ -522,7 +528,7 @@ def test_uniform_other(self, lower, upper, size, transform):
             initval=initval,
             transform=transform,
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=1)
+        self.check_vectortransform_elementwise_logp(model)
 
     @pytest.mark.parametrize(
         "mu,cov,size,shape",
@@ -536,7 +542,7 @@ def test_mvnormal_ordered(self, mu, cov, size, shape):
         model = self.build_model(
             pm.MvNormal, {"mu": mu, "cov": cov}, size=size, initval=initval, transform=tr.ordered
         )
-        self.check_vectortransform_elementwise_logp(model, vect_opt=1)
+        self.check_vectortransform_elementwise_logp(model)
 
 
 def test_triangular_transform():
