Use model context in model.logpt

Author: aseyboldt

pymc3/model.py

@@ -455,15 +455,17 @@ def dlogp_array(self):
     @memoize
     def logpt(self):
         """Theano scalar of log-probability of the model"""
-        factors = [var.logpt for var in self.basic_RVs] + self.potentials
-        return tt.add(*map(tt.sum, factors))
+        with self:
+            factors = [var.logpt for var in self.basic_RVs] + self.potentials
+            return tt.add(*map(tt.sum, factors))
 
     @property
     def varlogpt(self):
         """Theano scalar of log-probability of the unobserved random variables
            (excluding deterministic)."""
-        factors = [var.logpt for var in self.vars]
-        return tt.add(*map(tt.sum, factors))
+        with self:
+            factors = [var.logpt for var in self.vars]
+            return tt.add(*map(tt.sum, factors))
 
     @property
     def vars(self):
@@ -521,14 +523,16 @@ def Var(self, name, dist, data=None, total_size=None):
         name = self.name_for(name)
         if data is None:
             if getattr(dist, "transform", None) is None:
-                var = FreeRV(name=name, distribution=dist,
-                             total_size=total_size, model=self)
+                with self:
+                    var = FreeRV(name=name, distribution=dist,
+                                 total_size=total_size, model=self)
                 self.free_RVs.append(var)
             else:
-                var = TransformedRV(name=name, distribution=dist,
-                                    transform=dist.transform,
-                                    total_size=total_size,
-                                    model=self)
+                with self:
+                    var = TransformedRV(name=name, distribution=dist,
+                                        transform=dist.transform,
+                                        total_size=total_size,
+                                        model=self)
                 pm._log.debug('Applied {transform}-transform to {name}'
                               ' and added transformed {orig_name} to model.'.format(
                                 transform=dist.transform.name,
@@ -537,18 +541,20 @@ def Var(self, name, dist, data=None, total_size=None):
                 self.deterministics.append(var)
                 return var
         elif isinstance(data, dict):
-            var = MultiObservedRV(name=name, data=data, distribution=dist,
-                                  total_size=total_size, model=self)
+            with self:
+                var = MultiObservedRV(name=name, data=data, distribution=dist,
+                                      total_size=total_size, model=self)
             self.observed_RVs.append(var)
             if var.missing_values:
                 self.free_RVs += var.missing_values
                 self.missing_values += var.missing_values
                 for v in var.missing_values:
                     self.named_vars[v.name] = v
         else:
-            var = ObservedRV(name=name, data=data,
-                             distribution=dist,
-                             total_size=total_size, model=self)
+            with self:
+                var = ObservedRV(name=name, data=data,
+                                 distribution=dist,
+                                 total_size=total_size, model=self)
             self.observed_RVs.append(var)
             if var.missing_values:
                 self.free_RVs.append(var.missing_values)
@@ -615,11 +621,12 @@ def makefn(self, outs, mode=None, *args, **kwargs):
         -------
         Compiled Theano function
         """
-        return theano.function(self.vars, outs,
-                               allow_input_downcast=True,
-                               on_unused_input='ignore',
-                               accept_inplace=True,
-                               mode=mode, *args, **kwargs)
+        with self:
+            return theano.function(self.vars, outs,
+                                   allow_input_downcast=True,
+                                   on_unused_input='ignore',
+                                   accept_inplace=True,
+                                   mode=mode, *args, **kwargs)
 
     def fn(self, outs, mode=None, *args, **kwargs):
         """Compiles a Theano function which returns the values of `outs`