Add model.target_log_prob_fn()

pymc4/model/base.py

@@ -1,3 +1,4 @@
+import collections
 import biwrap
 import tensorflow as tf
 from tensorflow_probability import edward2 as ed
@@ -67,6 +68,14 @@ def get_mode(state, rv, *args, **kwargs):
             returns = self.session.run(list(values_collector.result.values()))
         return dict(zip(values_collector.result.keys(), returns))
 
+    def target_log_prob_fn(self, *args, **kwargs):
+        logp = 0
+        for i in self.unobserved.keys():
+            print(kwargs.get(i))
+            logp += self.unobserved[i].rv.distribution.log_prob(value=kwargs.get(i))
+
+        return logp
+
     def observe(self, **observations):
         self._observed = observations
         return self
@@ -83,6 +92,16 @@ def graph(self):
     def observed(self):
         return self._observed
 
+    @property
+    def unobserved(self):
+        unobserved = {}
+        for i in self.variables:
+            if self.variables[i] not in self.observed.values():
+                unobserved[i] = self.variables[i]
+
+        unobserved = collections.OrderedDict(unobserved)
+        return unobserved
+
     @property
     def variables(self):
         return self._variables

pymc4/util/interceptors.py

@@ -9,7 +9,7 @@
     'CollectLogProb'
 ]
 
-VariableDescription = collections.namedtuple('VariableDescription', 'Dist,shape')
+VariableDescription = collections.namedtuple('VariableDescription', 'Dist,shape,rv')
 
 
 class Interceptor(object):
@@ -85,7 +85,7 @@ def __init__(self):
     def after(self, rv, *args, **kwargs):
         name = kwargs["name"]
         if name not in self.result:
-            self.result[name] = VariableDescription(rv.distribution.__class__, rv.shape)
+            self.result[name] = VariableDescription(rv.distribution.__class__, rv.shape, rv)
         else:
             raise KeyError(name, 'Duplicate name')
         return rv