ensuring logpdf is run once
@@ -73,6 +73,8 @@ class KL_QP(object): # AbstractInfer):
model_trace = poutine.trace(
poutine.replay(self.model, guide_trace))(*args, **kwargs)
+ log_r = model_trace.log_pdf() - guide_trace.log_pdf()
+ log_r_per_sample.append(log_r)
model_traces.append(model_trace)
guide_traces.append(guide_trace)
@@ -80,23 +82,23 @@ class KL_QP(object): # AbstractInfer):
elbo = 0.0
for i in range(nr_particles):
elbo_particle = 0.0
+ log_r_s = 0.0
model_trace = model_traces[i]
guide_trace = guide_traces[i]
+ log_r = log_r_per_sample[i]
for name in model_trace.keys():
if model_trace[name]["type"] == "observe":
elbo_particle += model_trace[name]["log_pdf"]
elif model_trace[name]["type"] == "sample":
- if model_trace[name]["fn"].reparameterized:
- elbo_particle += model_trace[name]["log_pdf"]
- elbo_particle -= guide_trace[name]["log_pdf"]
- else:
elbo_particle += model_trace[name]["log_pdf"]
elbo_particle -= guide_trace[name]["log_pdf"]
+
else:
pass
elbo += elbo_particle
+
# gradients
loss = -elbo
@@ -124,6 +126,12 @@ class KL_QP(object): # AbstractInfer):
model_traces.append(model_trace)
guide_traces.append(guide_trace)
+ # guide_trace = poutine.trace(self.guide)(*args, **kwargs)
+ # model_trace = poutine.trace(
+ # poutine.replay(self.model, guide_trace))(*args, **kwargs)
+
+ # # compute losses
+ # log_r = model_trace.log_pdf() - guide_trace.log_pdf()
elbo = 0.0
for i in range(nr_particles):
