Make logpt work correctly for nested models and transforms

Author: brandonwillard

pymc3/distributions/__init__.py

@@ -26,7 +26,7 @@
 
 from aesara.tensor.random.op import RandomVariable
 
-from pymc3.distributions import _logcdf, _logp, logp_transform
+from pymc3.distributions import _logcdf, _logp
 
 if TYPE_CHECKING:
     from typing import Optional, Callable
@@ -111,12 +111,12 @@ def logp(op, value, *dist_params, **kwargs):
                 def logcdf(op, value, *dist_params, **kwargs):
                     return class_logcdf(value, *dist_params, **kwargs)
 
-            class_transform = clsdict.get("transform")
-            if class_transform:
-
-                @logp_transform.register(rv_type)
-                def transform(op, *args, **kwargs):
-                    return class_transform(*args, **kwargs)
+            # class_transform = clsdict.get("transform")
+            # if class_transform:
+            #
+            #     @logp_transform.register(rv_type)
+            #     def transform(op, *args, **kwargs):
+            #         return class_transform(*args, **kwargs)
 
             # Register the Aesara `RandomVariable` type as a subclass of this
             # `Distribution` type.
@@ -328,26 +328,17 @@ def _distr_parameters_for_repr(self):
 class Discrete(Distribution):
     """Base class for discrete distributions"""
 
-    def __init__(self, shape=(), dtype=None, defaults=("mode",), *args, **kwargs):
-        if dtype is None:
-            if aesara.config.floatX == "float32":
-                dtype = "int16"
-            else:
-                dtype = "int64"
-        if dtype != "int16" and dtype != "int64":
-            raise TypeError("Discrete classes expect dtype to be int16 or int64.")
+    def __new__(cls, name, *args, **kwargs):
 
-        super().__init__(shape, dtype, defaults=defaults, *args, **kwargs)
+        if kwargs.get("transform", None):
+            raise ValueError("Transformations for discrete distributions")
+
+        return super().__new__(cls, name, *args, **kwargs)
 
 
 class Continuous(Distribution):
     """Base class for continuous distributions"""
 
-    def __init__(self, shape=(), dtype=None, defaults=("median", "mean", "mode"), *args, **kwargs):
-        if dtype is None:
-            dtype = aesara.config.floatX
-        super().__init__(shape, dtype, defaults=defaults, *args, **kwargs)
-
 
 class DensityDist(Distribution):
     """Distribution based on a given log density function.

pymc3/distributions/distribution.py

@@ -388,7 +388,7 @@ class Dirichlet(Continuous):
     rv_op = dirichlet
 
     @classmethod
-    def dist(cls, a, **kwargs):
+    def dist(cls, a, transform=transforms.stick_breaking, **kwargs):
 
         a = at.as_tensor_variable(a)
         # mean = a / at.sum(a)
@@ -419,15 +419,6 @@ def logp(value, a):
             broadcast_conditions=False,
         )
 
-    def transform(rv_var):
-
-        if rv_var.ndim == 1 or rv_var.broadcastable[-1]:
-            # If this variable is just a bunch of scalars/degenerate
-            # Dirichlets, we can't transform it
-            return None
-
-        return transforms.stick_breaking
-
     def _distr_parameters_for_repr(self):
         return ["a"]
 

pymc3/distributions/multivariate.py

@@ -218,7 +218,7 @@ def jacobian_det(self, rv_var, rv_value):
             s = at.nnet.softplus(-rv_value)
             return at.log(b - a) - 2 * s - rv_value
         else:
-            return rv_value
+            return at.ones_like(rv_value)
 
 
 interval = Interval
@@ -286,6 +286,11 @@ class StickBreaking(Transform):
     name = "stickbreaking"
 
     def forward(self, rv_var, rv_value):
+        if rv_var.ndim == 1 or rv_var.broadcastable[-1]:
+            # If this variable is just a bunch of scalars/degenerate
+            # Dirichlets, we can't transform it
+            return rv_value
+
         x = rv_value.T
         n = x.shape[0]
         lx = at.log(x)
@@ -294,6 +299,11 @@ def forward(self, rv_var, rv_value):
         return floatX(y.T)
 
     def backward(self, rv_var, rv_value):
+        if rv_var.ndim == 1 or rv_var.broadcastable[-1]:
+            # If this variable is just a bunch of scalars/degenerate
+            # Dirichlets, we can't transform it
+            return rv_value
+
         y = rv_value.T
         y = at.concatenate([y, -at.sum(y, 0, keepdims=True)])
         # "softmax" with vector support and no deprication warning:
@@ -302,6 +312,11 @@ def backward(self, rv_var, rv_value):
         return floatX(x.T)
 
     def jacobian_det(self, rv_var, rv_value):
+        if rv_var.ndim == 1 or rv_var.broadcastable[-1]:
+            # If this variable is just a bunch of scalars/degenerate
+            # Dirichlets, we can't transform it
+            return at.ones_like(rv_value)
+
         y = rv_value.T
         Km1 = y.shape[0] + 1
         sy = at.sum(y, 0, keepdims=True)

pymc3/distributions/transforms.py

@@ -239,6 +239,8 @@ def build_named_node_tree(graphs):
 
 T = TypeVar("T", bound="ContextMeta")
 
+no_transform_object = object()
+
 
 class ContextMeta(type):
     """Functionality for objects that put themselves in a context using
@@ -1047,7 +1049,9 @@ def add_coords(self, coords):
             else:
                 self.coords[name] = coords[name]
 
-    def register_rv(self, rv_var, name, data=None, total_size=None, dims=None, transform=None):
+    def register_rv(
+        self, rv_var, name, data=None, total_size=None, dims=None, transform=no_transform_object
+    ):
         """Register an (un)observed random variable with the model.
 
         Parameters
@@ -1104,24 +1108,21 @@ def register_rv(self, rv_var, name, data=None, total_size=None, dims=None, trans
         if aesara.config.compute_test_value != "off":
             value_var.tag.test_value = rv_var.tag.test_value
 
-        value_var.name = f"{rv_var.name}_value"
+        value_var.name = rv_var.name
 
         rv_var.tag.value_var = value_var
 
         # Make the value variable a transformed value variable,
         # if there's an applicable transform
-        transform = transform or logp_transform(rv_var.owner.op)
+        if transform is no_transform_object:
+            transform = logp_transform(rv_var.owner.op)
 
         if transform is not None:
             value_var.tag.transform = transform
             value_var.name = f"{value_var.name}_{transform.name}__"
             if aesara.config.compute_test_value != "off":
                 value_var.tag.test_value = transform.forward(rv_var, value_var).tag.test_value
 
-            # The transformed variable needs to be a named variable in the
-            # model, too
-            self.named_vars[value_var.name] = value_var
-
         self.add_random_variable(rv_var, dims)
 
         return rv_var
@@ -1173,7 +1174,7 @@ def __getitem__(self, key):
             except KeyError:
                 raise e
 
-    def makefn(self, outs, mode=None, transformed=True, *args, **kwargs):
+    def makefn(self, outs, mode=None, *args, **kwargs):
         """Compiles a Aesara function which returns ``outs`` and takes the variable
         ancestors of ``outs`` as inputs.
 
@@ -1187,11 +1188,8 @@ def makefn(self, outs, mode=None, transformed=True, *args, **kwargs):
         Compiled Aesara function
         """
         with self:
-            vars = [
-                v if not transformed else getattr(v.tag, "transformed_var", v) for v in self.vars
-            ]
             return aesara.function(
-                vars,
+                self.vars,
                 outs,
                 allow_input_downcast=True,
                 on_unused_input="ignore",
@@ -1324,7 +1322,10 @@ def check_test_point(self, test_point=None, round_vals=2):
 
         return Series(
             {
-                rv.name: np.round(self.fn(logpt_sum(rv))(test_point), round_vals)
+                rv.name: np.round(
+                    self.fn(logpt_sum(rv, getattr(rv.tag, "observations", None)))(test_point),
+                    round_vals,
+                )
                 for rv in self.basic_RVs
             },
             name="Log-probability of test_point",

pymc3/model.py

@@ -1723,7 +1723,9 @@ def sample_posterior_predictive(
         return {}
 
     if not hasattr(_trace, "varnames"):
-        inputs_and_names = [(i, i.name) for i in rv_ancestors(vars_to_sample)]
+        inputs_and_names = [
+            (rv, rv.name) for rv in rv_ancestors(vars_to_sample, walk_past_rvs=True)
+        ]
         inputs, input_names = zip(*inputs_and_names)
     else:
         input_names = _trace.varnames

pymc3/sampling.py

@@ -19,7 +19,7 @@
 """
 from warnings import warn
 
-import aesara.tensor as aet
+import aesara.tensor as at
 
 from aesara.graph.basic import graph_inputs
 from numpy import arange, array, cumsum, empty, exp, max, nested_iters, searchsorted
@@ -81,7 +81,7 @@ def elemwise_logp(model, var):
         v_logp = logpt(v)
         if var in graph_inputs([v_logp]):
             terms.append(v_logp)
-    return model.fn(aet.add(*terms))
+    return model.fn(at.add(*terms))
 
 
 def categorical(prob, shape):