From 4efb780c46c27b435934985a635649c4e84616f7 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Tue, 2 Jan 2018 21:41:07 -0500
Subject: [PATCH 01/42] fix typos in docstring

---
 edward/inferences/klpq.py                  | 2 +-
 edward/inferences/variational_inference.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/edward/inferences/klpq.py b/edward/inferences/klpq.py
index 0dcaa71b4..dcf4dd9bf 100644
--- a/edward/inferences/klpq.py
+++ b/edward/inferences/klpq.py
@@ -32,7 +32,7 @@ class KLpq(VariationalInference):
 
   with respect to $\\theta$.
 
-  In conditional inference, we infer $z` in $p(z, \\beta
+  In conditional inference, we infer $z$ in $p(z, \\beta
   \mid x)$ while fixing inference over $\\beta$ using another
   distribution $q(\\beta)$. During gradient calculation, instead
   of using the model's density
diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index 411a16b48..171eca56b 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -18,7 +18,7 @@ class VariationalInference(Inference):
   variational inference methods inherit from `VariationalInference`,
   sharing methods such as a default optimizer.
 
-  To build an algorithm inheriting from `VariaitonalInference`, one
+  To build an algorithm inheriting from `VariationalInference`, one
   must at the minimum implement `build_loss_and_gradients`: it
   determines the loss function and gradients to apply for a given
   optimizer.

From 7e43d1b2ff81ca94af347a78d9a39cda89c437ef Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 6 Jan 2018 19:23:49 -0500
Subject: [PATCH 02/42] add multinomial-dirichlet test, empty
 `RejectionSamplingKLqp` class

---
 edward/__init__.py            |  3 ++-
 edward/inferences/__init__.py |  1 +
 edward/inferences/klqp.py     |  9 +++++++++
 tests/inferences/test_klqp.py | 21 ++++++++++++++++++++-
 4 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/edward/__init__.py b/edward/__init__.py
index 4997a89f3..4a2d0df35 100644
--- a/edward/__init__.py
+++ b/edward/__init__.py
@@ -14,7 +14,7 @@
     HMC, MetropolisHastings, SGLD, SGHMC, \
     KLpq, KLqp, ReparameterizationKLqp, ReparameterizationKLKLqp, \
     ReparameterizationEntropyKLqp, ScoreKLqp, ScoreKLKLqp, ScoreEntropyKLqp, \
-    ScoreRBKLqp, WakeSleep, GANInference, BiGANInference, WGANInference, \
+    ScoreRBKLqp, RejectionSamplingKLqp, WakeSleep, GANInference, BiGANInference, WGANInference, \
     ImplicitKLqp, MAP, Laplace, complete_conditional, Gibbs
 from edward.models import RandomVariable
 from edward.util import check_data, check_latent_vars, copy, dot, \
@@ -52,6 +52,7 @@
     'ScoreKLKLqp',
     'ScoreEntropyKLqp',
     'ScoreRBKLqp',
+    'RejectionSamplingKLqp',
     'WakeSleep',
     'GANInference',
     'BiGANInference',
diff --git a/edward/inferences/__init__.py b/edward/inferences/__init__.py
index 38262fcb7..fd22d9e8c 100644
--- a/edward/inferences/__init__.py
+++ b/edward/inferences/__init__.py
@@ -42,6 +42,7 @@
     'ScoreKLKLqp',
     'ScoreEntropyKLqp',
     'ScoreRBKLqp',
+    'RejectionSamplingKLqp',
     'Laplace',
     'MAP',
     'MetropolisHastings',
diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index c6a9b386b..c5e14ffc7 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -592,6 +592,15 @@ def build_loss_and_gradients(self, var_list):
     return build_score_rb_loss_and_gradients(self, var_list)
 
 
+class RejectionSamplingKLqp(VariationalInference):
+
+    """
+    """
+
+    def build_loss_and_gradients(self, var_list):
+        return
+
+
 def build_reparam_loss_and_gradients(inference, var_list):
   """Build loss function. Its automatic differentiation
   is a stochastic gradient of
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index baf729394..4e76a0f0a 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -6,7 +6,7 @@
 import numpy as np
 import tensorflow as tf
 
-from edward.models import Bernoulli, Normal
+from edward.models import Bernoulli, Normal, Dirichlet, Multinomial
 
 
 class test_klqp_class(tf.test.TestCase):
@@ -43,6 +43,20 @@ def _test_normal_normal(self, Inference, default, *args, **kwargs):
       self.assertEqual(new_t, 0)
       self.assertNotEqual(old_variables, new_variables)
 
+  def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
+      x_data = tf.constant([2, 7, 1], dtype=np.float32)
+
+      probs = Dirichlet([1., 1., 1.])
+      x = Multinomial(total_count=10.0, probs=probs, sample_shape=1)
+
+      qalpha = tf.Variable(tf.random_normal([3]))
+      qprobs = Dirichlet(qalpha)
+
+      # analytic solution: Dirichlet(alpha=[1+2, 1+7, 1+1])
+      inference = Inference({probs: qprobs}, data={x: x_data})
+
+      inference.run(*args, **kwargs)
+
   def _test_model_parameter(self, Inference, *args, **kwargs):
     with self.test_session() as sess:
       x_data = np.array([0, 1, 0, 0, 0, 0, 0, 0, 0, 1])
@@ -109,6 +123,11 @@ def test_score_rb_klqp(self):
         ed.ScoreRBKLqp, default=True, n_samples=5, n_iter=5000)
     self._test_model_parameter(ed.ScoreRBKLqp, n_iter=50)
 
+  def test_rejection_sampling_klqp(self):
+    self._test_multinomial_dirichlet(
+      ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
+
+
 if __name__ == '__main__':
   ed.set_seed(42)
   tf.test.main()

From 7a5f90e2a03f7a4bc79b4a32996bcc2df88791d4 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Thu, 11 Jan 2018 22:19:30 -0500
Subject: [PATCH 03/42] remove `sample_shape=1`

---
 tests/inferences/test_klqp.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 4e76a0f0a..edace0f5d 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -47,7 +47,7 @@ def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
       x_data = tf.constant([2, 7, 1], dtype=np.float32)
 
       probs = Dirichlet([1., 1., 1.])
-      x = Multinomial(total_count=10.0, probs=probs, sample_shape=1)
+      x = Multinomial(total_count=10.0, probs=probs)
 
       qalpha = tf.Variable(tf.random_normal([3]))
       qprobs = Dirichlet(qalpha)

From 94a1bc39936d71d6bb26c224a713b3563db94ea7 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sun, 14 Jan 2018 11:48:45 -0500
Subject: [PATCH 04/42] add poisson-gamma test

---
 tests/inferences/test_klqp.py | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index edace0f5d..a9338226a 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -6,7 +6,8 @@
 import numpy as np
 import tensorflow as tf
 
-from edward.models import Bernoulli, Normal, Dirichlet, Multinomial
+from edward.models import Bernoulli, Normal, Dirichlet, Multinomial, \
+  Gamma, Poisson
 
 
 class test_klqp_class(tf.test.TestCase):
@@ -43,6 +44,23 @@ def _test_normal_normal(self, Inference, default, *args, **kwargs):
       self.assertEqual(new_t, 0)
       self.assertNotEqual(old_variables, new_variables)
 
+  def _test_poisson_gamma(self, Inference, *args, **kwargs):
+      x_data = np.array([2, 8, 3, 6, 1], dtype=np.int32)
+
+      rate = Gamma(5.0, 1.0)
+      x = Poisson(rate=rate, sample_shape=5)
+
+      qalpha = tf.Variable(tf.random_normal([]))
+      qbeta = tf.Variable(tf.random_normal([]))
+      qgamma = Gamma(qalpha, qbeta)
+
+      # sum(x_data) = 20
+      # len(x_data) = 5
+      # analytic solution: Gamma(alpha=5+20, beta=1+5)
+      inference = Inference({rate: qgamma}, data={x: x_data})
+
+      inference.run(*args, **kwargs)
+
   def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
       x_data = tf.constant([2, 7, 1], dtype=np.float32)
 
@@ -124,8 +142,10 @@ def test_score_rb_klqp(self):
     self._test_model_parameter(ed.ScoreRBKLqp, n_iter=50)
 
   def test_rejection_sampling_klqp(self):
-    self._test_multinomial_dirichlet(
-      ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
+    self._test_poisson_gamma(
+      ed.ReparameterizationKLqp, n_samples=5, n_iter=5000)
+    # self._test_multinomial_dirichlet(
+    #   ed.ReparameterizationKLqp, n_samples=5, n_iter=5000)
 
 
 if __name__ == '__main__':

From a4c87cc01f9c77be7b173a73913edffd7dd2f9fc Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sun, 14 Jan 2018 21:24:44 -0500
Subject: [PATCH 05/42] WIP: begin to implement RSVI logic

---
 edward/inferences/klqp.py     | 124 +++++++++++++++++++++++++++++++++-
 tests/inferences/test_klqp.py |  35 ++++++++--
 2 files changed, 154 insertions(+), 5 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 8aff267f4..45b6baa65 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -616,13 +616,57 @@ def build_loss_and_gradients(self, var_list):
     return build_score_rb_loss_and_gradients(self, var_list)
 
 
+# TODO: you can probably make another base class?
 class RejectionSamplingKLqp(VariationalInference):
 
     """
     """
 
+    def __init__(self, latent_vars=None, data=None, rejection_sampler_vars=None):
+      """Create an inference algorithm.
+
+      # TODO: update me
+
+      Args:
+        latent_vars: list of RandomVariable or
+                     dict of RandomVariable to RandomVariable.
+          Collection of random variables to perform inference on. If
+          list, each random variable will be implictly optimized using a
+          `Normal` random variable that is defined internally with a
+          free parameter per location and scale and is initialized using
+          standard normal draws. The random variables to approximate
+          must be continuous.
+      """
+      super(RejectionSamplingKLqp, self).__init__(latent_vars, data)
+      self.rejection_sampler_vars = rejection_sampler_vars
+
+    def initialize(self, n_samples=1, *args, **kwargs):
+      """Initialize inference algorithm. It initializes hyperparameters
+      and builds ops for the algorithm's computation graph.
+
+      Args:
+        n_samples: int, optional.
+          Number of samples from variational model for calculating
+          stochastic gradients.
+      """
+      self.n_samples = n_samples
+      return super(RejectionSamplingKLqp, self).initialize(*args, **kwargs)
+
     def build_loss_and_gradients(self, var_list):
-        return
+      return build_rejection_sampling_loss_and_gradients(self, var_list)
+
+    # TODO: pass in the real `qalpha` and `qbeta`
+    # TODO: pass in the real `qz`
+    def sample(self, qz, alpha=5, beta=5):
+      qz_class = qz.__class__
+      while True:
+        epsilon = self.rejection_sampler_vars[qz_class]['epsilon_likelihood'].value()
+        z = self.rejection_sampler_vars[qz_class]['reparam_func'](epsilon)
+        eps_prob = self.rejection_sampler_vars[qz_class]['epsilon_likelihood'].prob(epsilon)
+        qz_prob = qz.prob(z)
+        random_uniform = tf.random_uniform([])
+        if random_uniform * self.rejection_sampler_vars[qz_class]['m'] * eps_prob <= qz_prob:
+          return epsilon
 
 
 def build_reparam_loss_and_gradients(inference, var_list):
@@ -1136,3 +1180,81 @@ def build_score_rb_loss_and_gradients(inference, var_list):
   grads_vars.extend(model_vars)
   grads_and_vars = list(zip(grads, grads_vars))
   return loss, grads_and_vars
+
+
+def build_rejection_sampling_loss_and_gradients(inference, var_list):
+    """
+    """
+    p_log_prob = [0.0] * inference.n_samples
+    q_log_prob = [0.0] * inference.n_samples
+    r_log_prob = [0.0] * inference.n_samples
+    base_scope = tf.get_default_graph().unique_name("inference") + '/'
+    for s in range(inference.n_samples):
+      # Form dictionary in order to replace conditioning on prior or
+      # observed variable with conditioning on a specific value.
+      scope = base_scope + tf.get_default_graph().unique_name("sample")
+      dict_swap = {}
+      for x, qx in six.iteritems(inference.data):
+        if isinstance(x, RandomVariable):
+          if isinstance(qx, RandomVariable):
+            qx_copy = copy(qx, scope=scope)
+            dict_swap[x] = qx_copy.value()
+          else:
+            dict_swap[x] = qx
+
+      for z, qz in six.iteritems(inference.latent_vars):
+        # Copy q(z) to obtain new set of posterior samples.
+        # You need to check if the things need to be RSVI'd in the first place
+        qz_class = qz.__class__
+        epsilon_likelihood = self.rejection_sampler_vars[qz_class]['epsilon_likelihood']
+        qz_copy = copy(qz, scope=scope)
+
+        if 'rsvi':
+          # --- RSVI
+          epsilon = self.sample(qz_copy)
+          z = self.rejection_sampler_vars[qz_class]['reparam_func'](epsilon)
+          # RSVI ---
+        else:
+          z = qz_copy.value()
+        dict_swap[z] = z
+
+        q_log_prob[s] += tf.reduce_sum(
+            inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
+        r_log_prob[s] += tf.reduce_sum(
+            inference.scale.get(z, 1.0) * epsilon_likelihood.log_prob(dict_swap[z]))
+
+      for z in six.iterkeys(inference.latent_vars):
+        z_copy = copy(z, dict_swap, scope=scope)
+        p_log_prob[s] += tf.reduce_sum(
+            inference.scale.get(z, 1.0) * z_copy.log_prob(dict_swap[z]))
+
+      for x in six.iterkeys(inference.data):
+        if isinstance(x, RandomVariable):
+          x_copy = copy(x, dict_swap, scope=scope)
+          p_log_prob[s] += tf.reduce_sum(
+              inference.scale.get(x, 1.0) * x_copy.log_prob(dict_swap[x]))
+
+    p_log_prob = tf.reduce_mean(p_log_prob)
+    q_log_prob = tf.reduce_mean(q_log_prob)
+    r_log_prob = tf.reduce_mean(r_log_prob)
+
+    q_entropy = tf.reduce_sum([
+        tf.reduce_sum(qz.entropy())
+        for z, qz in six.iteritems(inference.latent_vars)])
+
+    reg_penalty = tf.reduce_sum(tf.losses.get_regularization_losses())
+
+    if inference.logging:
+      tf.summary.scalar("loss/p_log_prob", p_log_prob,
+                        collections=[inference._summary_key])
+      tf.summary.scalar("loss/q_entropy", q_entropy,
+                        collections=[inference._summary_key])
+      tf.summary.scalar("loss/reg_penalty", reg_penalty,
+                        collections=[inference._summary_key])
+
+    loss = -(p_log_prob + q_entropy - reg_penalty)
+
+    # Here you have to carefully return the gradients; check what we do elsewhere
+    grads = tf.gradients(loss, var_list)
+    grads_and_vars = list(zip(grads, var_list))
+    return loss, grads_and_vars
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index a9338226a..be0e8fe3e 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -50,14 +50,41 @@ def _test_poisson_gamma(self, Inference, *args, **kwargs):
       rate = Gamma(5.0, 1.0)
       x = Poisson(rate=rate, sample_shape=5)
 
-      qalpha = tf.Variable(tf.random_normal([]))
-      qbeta = tf.Variable(tf.random_normal([]))
+      qalpha = tf.nn.softplus(tf.Variable(tf.random_normal([])))
+      qbeta = tf.nn.softplus(tf.Variable(tf.random_normal([])))
       qgamma = Gamma(qalpha, qbeta)
 
+      # Gamma rejection sampler variables
+      def gamma_reparam_func(epsilon, alpha, beta):
+        def _gamma_reparam_func(alpha, beta, epsilon=epsilon)
+          a = alpha - (1. / 3)
+          b = np.sqrt(9 * alpha - 3)
+          c = 1 + (epsilon / b)
+          z = a * c**3
+          return z
+
+        if alpha < 1:
+          z_tilde = _gamma_reparam_func(alpha=alpha + 1, beta=beta)
+          u = np.random.uniform()
+          z = u ** (1 / alpha) * z_tilde
+        else:
+          z = _gamma_reparam_func(alpha=alpha, beta=beta)
+        if beta != 1:
+          z /= beta
+
+        return z
+
+      gamma_rejection_sampler_vars = {
+        'reparam_func': gamma_reparam_func,
+        'epsilon_likelihood': Normal(loc=0.0, scale=1.0),
+        'm': 10.
+      }
+
       # sum(x_data) = 20
       # len(x_data) = 5
       # analytic solution: Gamma(alpha=5+20, beta=1+5)
-      inference = Inference({rate: qgamma}, data={x: x_data})
+      inference = Inference({rate: qgamma}, data={x: x_data},
+          rejection_sampler_vars={Gamma: gamma_rejection_sampler_vars})
 
       inference.run(*args, **kwargs)
 
@@ -143,7 +170,7 @@ def test_score_rb_klqp(self):
 
   def test_rejection_sampling_klqp(self):
     self._test_poisson_gamma(
-      ed.ReparameterizationKLqp, n_samples=5, n_iter=5000)
+      ed.RejectionSamplingKLqp, n_samples=1, n_iter=5000)
     # self._test_multinomial_dirichlet(
     #   ed.ReparameterizationKLqp, n_samples=5, n_iter=5000)
 

From 163414c2731680bac81358440216a5399ae0ee98 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 15 Jan 2018 14:40:46 -0500
Subject: [PATCH 06/42] WIP: implement RSVI gradients

---
 edward/inferences/klqp.py                  |  56 ++++++-----
 edward/inferences/variational_inference.py |   3 +
 tests/inferences/test_klqp.py              | 104 +++++++++++----------
 3 files changed, 91 insertions(+), 72 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 45b6baa65..881124903 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -616,7 +616,7 @@ def build_loss_and_gradients(self, var_list):
     return build_score_rb_loss_and_gradients(self, var_list)
 
 
-# TODO: you can probably make another base class?
+# TODO: you can probably make another base class that implements a `sample` method?
 class RejectionSamplingKLqp(VariationalInference):
 
     """
@@ -655,19 +655,6 @@ def initialize(self, n_samples=1, *args, **kwargs):
     def build_loss_and_gradients(self, var_list):
       return build_rejection_sampling_loss_and_gradients(self, var_list)
 
-    # TODO: pass in the real `qalpha` and `qbeta`
-    # TODO: pass in the real `qz`
-    def sample(self, qz, alpha=5, beta=5):
-      qz_class = qz.__class__
-      while True:
-        epsilon = self.rejection_sampler_vars[qz_class]['epsilon_likelihood'].value()
-        z = self.rejection_sampler_vars[qz_class]['reparam_func'](epsilon)
-        eps_prob = self.rejection_sampler_vars[qz_class]['epsilon_likelihood'].prob(epsilon)
-        qz_prob = qz.prob(z)
-        random_uniform = tf.random_uniform([])
-        if random_uniform * self.rejection_sampler_vars[qz_class]['m'] * eps_prob <= qz_prob:
-          return epsilon
-
 
 def build_reparam_loss_and_gradients(inference, var_list):
   """Build loss function. Its automatic differentiation
@@ -1204,19 +1191,38 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
 
       for z, qz in six.iteritems(inference.latent_vars):
         # Copy q(z) to obtain new set of posterior samples.
-        # You need to check if the things need to be RSVI'd in the first place
-        qz_class = qz.__class__
-        epsilon_likelihood = self.rejection_sampler_vars[qz_class]['epsilon_likelihood']
         qz_copy = copy(qz, scope=scope)
 
+        # Of course, this will evaluate to `True`. We just do this as a simple first pass.
         if 'rsvi':
           # --- RSVI
-          epsilon = self.sample(qz_copy)
-          z = self.rejection_sampler_vars[qz_class]['reparam_func'](epsilon)
+
+          # Get variable shortnames
+          qz_class = qz.__class__
+          epsilon_likelihood = inference.rejection_sampler_vars[qz_class]['epsilon_likelihood']
+          reparam_func = inference.rejection_sampler_vars[qz_class]['reparam_func']
+          m = inference.rejection_sampler_vars[qz_class]['m']
+          alpha = qz.parameters['concentration']
+          beta = qz.parameters['rate']
+
+          # Sample
+
+          # TODO: pass in the real `qalpha` and `qbeta`
+          # TODO: pass in the real `qz`
+          epsilon = epsilon_likelihood.value()
+          sample = reparam_func(epsilon, alpha, beta)
+          eps_prob = epsilon_likelihood.prob(epsilon)
+          qz_prob = qz.prob(sample)
+          random_uniform = tf.random_uniform([])
+
+          # We need this line. However, let's just accept for now.
+          # if random_uniform * m * eps_prob <= qz_prob:
+
           # RSVI ---
         else:
           z = qz_copy.value()
-        dict_swap[z] = z
+
+        dict_swap[z] = sample
 
         q_log_prob[s] += tf.reduce_sum(
             inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
@@ -1254,7 +1260,13 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
 
     loss = -(p_log_prob + q_entropy - reg_penalty)
 
-    # Here you have to carefully return the gradients; check what we do elsewhere
-    grads = tf.gradients(loss, var_list)
+    # RSVI gradient components
+    model_grad = tf.gradients(p_log_prob, sample)[0]
+    q_entropy_grad = tf.gradients(q_entropy, var_list)
+    g_rep = [model_grad * grad for grad in tf.gradients(sample, var_list)]
+    g_cor = [p_log_prob * grad for grad in tf.gradients(q_log_prob - r_log_prob, var_list)]
+    grad_summands = zip(*[g_rep, g_cor, q_entropy_grad])
+
+    grads = [tf.reduce_sum(summand) for summand in grad_summands]
     grads_and_vars = list(zip(grads, var_list))
     return loss, grads_and_vars
diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index 171eca56b..633094a3f 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -117,6 +117,7 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
 
     with tf.variable_scope(None, default_name="optimizer") as scope:
       if not use_prettytensor:
+        self.grads_and_vars = grads_and_vars  # debug
         self.train = optimizer.apply_gradients(grads_and_vars,
                                                global_step=global_step)
       else:
@@ -151,6 +152,8 @@ def update(self, feed_dict=None):
         feed_dict[key] = value
 
     sess = get_session()
+    # TODO: delete me
+    # grads_and_vars_debug = sess.run([self.grads_and_vars], feed_dict)
     _, t, loss = sess.run([self.train, self.increment_t, self.loss], feed_dict)
 
     if self.debug:
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index be0e8fe3e..185f6aa35 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -45,62 +45,66 @@ def _test_normal_normal(self, Inference, default, *args, **kwargs):
       self.assertNotEqual(old_variables, new_variables)
 
   def _test_poisson_gamma(self, Inference, *args, **kwargs):
-      x_data = np.array([2, 8, 3, 6, 1], dtype=np.int32)
-
-      rate = Gamma(5.0, 1.0)
-      x = Poisson(rate=rate, sample_shape=5)
-
-      qalpha = tf.nn.softplus(tf.Variable(tf.random_normal([])))
-      qbeta = tf.nn.softplus(tf.Variable(tf.random_normal([])))
-      qgamma = Gamma(qalpha, qbeta)
-
-      # Gamma rejection sampler variables
-      def gamma_reparam_func(epsilon, alpha, beta):
-        def _gamma_reparam_func(alpha, beta, epsilon=epsilon)
-          a = alpha - (1. / 3)
-          b = np.sqrt(9 * alpha - 3)
-          c = 1 + (epsilon / b)
-          z = a * c**3
-          return z
-
-        if alpha < 1:
-          z_tilde = _gamma_reparam_func(alpha=alpha + 1, beta=beta)
-          u = np.random.uniform()
-          z = u ** (1 / alpha) * z_tilde
-        else:
-          z = _gamma_reparam_func(alpha=alpha, beta=beta)
-        if beta != 1:
-          z /= beta
-
-        return z
-
-      gamma_rejection_sampler_vars = {
-        'reparam_func': gamma_reparam_func,
-        'epsilon_likelihood': Normal(loc=0.0, scale=1.0),
-        'm': 10.
-      }
-
-      # sum(x_data) = 20
-      # len(x_data) = 5
-      # analytic solution: Gamma(alpha=5+20, beta=1+5)
-      inference = Inference({rate: qgamma}, data={x: x_data},
-          rejection_sampler_vars={Gamma: gamma_rejection_sampler_vars})
+      with self.test_session() as sess:
+          x_data = np.array([2, 8, 3, 6, 1], dtype=np.int32)
 
-      inference.run(*args, **kwargs)
+          rate = Gamma(5.0, 1.0)
+          x = Poisson(rate=rate, sample_shape=5)
+
+          qalpha = tf.nn.softplus(tf.Variable(tf.random_normal([])))
+          qbeta = tf.nn.softplus(tf.Variable(tf.random_normal([])))
+          qgamma = Gamma(qalpha, qbeta)
+
+          # Gamma rejection sampler variables
+          def gamma_reparam_func(epsilon, alpha, beta):
+
+            def _gamma_reparam_func(alpha=alpha, beta=beta, epsilon=epsilon):
+              a = alpha - (1. / 3)
+              b = tf.sqrt(9 * alpha - 3)
+              c = 1 + (epsilon / b)
+              z = a * c**3
+              return z
+
+            def _gamma_reparam_func_alpha_lt_1(alpha=alpha, beta=beta, epsilon=epsilon):
+              z_tilde = _gamma_reparam_func(alpha=alpha + 1, beta=beta)
+              u = np.random.uniform()
+              z = u ** (1 / alpha) * z_tilde
+              return z
+
+            z = tf.cond(tf.less(alpha, 1.), _gamma_reparam_func_alpha_lt_1, _gamma_reparam_func)
+            z = tf.cond(tf.equal(beta, 1.), lambda: z, lambda: tf.divide(z, beta))
+            return z
+
+          gamma_rejection_sampler_vars = {
+            'reparam_func': gamma_reparam_func,
+            'epsilon_likelihood': Normal(loc=0.0, scale=1.0),
+            'm': 10.
+          }
+
+          # sum(x_data) = 20
+          # len(x_data) = 5
+          # analytic solution: Gamma(alpha=5+20, beta=1+5)
+          inference = Inference({rate: qgamma}, data={x: x_data},
+              rejection_sampler_vars={Gamma: gamma_rejection_sampler_vars})
+
+          inference.run(*args, **kwargs)
+
+          import pdb; pdb.set_trace()
 
   def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
-      x_data = tf.constant([2, 7, 1], dtype=np.float32)
+      with self.test_session() as sess:
+          x_data = tf.constant([2, 7, 1], dtype=np.float32)
 
-      probs = Dirichlet([1., 1., 1.])
-      x = Multinomial(total_count=10.0, probs=probs)
+          probs = Dirichlet([1., 1., 1.])
+          x = Multinomial(total_count=10.0, probs=probs)
 
-      qalpha = tf.Variable(tf.random_normal([3]))
-      qprobs = Dirichlet(qalpha)
+          qalpha = tf.Variable(tf.random_normal([3]))
+          qprobs = Dirichlet(qalpha)
 
-      # analytic solution: Dirichlet(alpha=[1+2, 1+7, 1+1])
-      inference = Inference({probs: qprobs}, data={x: x_data})
+          # analytic solution: Dirichlet(alpha=[1+2, 1+7, 1+1])
+          inference = Inference({probs: qprobs}, data={x: x_data})
 
-      inference.run(*args, **kwargs)
+          inference.run(*args, **kwargs)
 
   def _test_model_parameter(self, Inference, *args, **kwargs):
     with self.test_session() as sess:
@@ -172,7 +176,7 @@ def test_rejection_sampling_klqp(self):
     self._test_poisson_gamma(
       ed.RejectionSamplingKLqp, n_samples=1, n_iter=5000)
     # self._test_multinomial_dirichlet(
-    #   ed.ReparameterizationKLqp, n_samples=5, n_iter=5000)
+    #   ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
 
 
 if __name__ == '__main__':

From f162135901a07d5f9a77cd880e528e5d8f37c876 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Thu, 18 Jan 2018 23:40:39 -0500
Subject: [PATCH 07/42] add scrap notebook with gradient update algo

---
 scrap.ipynb | 227 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 227 insertions(+)
 create mode 100644 scrap.ipynb

diff --git a/scrap.ipynb b/scrap.ipynb
new file mode 100644
index 000000000..2a6ddb2bc
--- /dev/null
+++ b/scrap.ipynb
@@ -0,0 +1,227 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "from edward.models import Normal\n",
+    "import tensorflow as tf"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):\n",
+    "    t = 0.1\n",
+    "    delta = 10e-3\n",
+    "    eta = 1e-1\n",
+    "\n",
+    "    update_ops = []\n",
+    "    for i, (grad, var) in enumerate(grads_and_vars):\n",
+    "        update_s_n_op = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )\n",
+    "        update_ops.append(update_s_n_op)\n",
+    "        \n",
+    "        p_n_first = eta * n**(-.5 + delta)\n",
+    "        p_n_second = 1 / (1 + tf.sqrt(s_n[i]))\n",
+    "        p_n = p_n_first * p_n_second\n",
+    "        \n",
+    "        apply_grad_op = var.assign_add(-p_n * grad)\n",
+    "        update_ops.append(apply_grad_op)\n",
+    "    return update_ops"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {
+    "collapsed": false,
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "w1 = tf.Variable(tf.random_normal([]))\n",
+    "w2 = tf.Variable(tf.random_normal([]))\n",
+    "var_list = [w1, w2]\n",
+    "\n",
+    "x = tf.constant([3., 4., 5.])\n",
+    "y = tf.constant([.8, .1, .1])\n",
+    "\n",
+    "pred = tf.nn.softmax(x * w1 * w2)\n",
+    "loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))\n",
+    "grads = tf.gradients(loss, var_list)\n",
+    "grads_and_vars = list(zip(grads, var_list))\n",
+    "\n",
+    "s_n = tf.Variable(tf.zeros(2))\n",
+    "n = tf.Variable(tf.constant(1.))\n",
+    "\n",
+    "train = alp_optimizer_apply_gradients(n, s_n, grads_and_vars)\n",
+    "increment_n = n.assign_add(1.)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {
+    "collapsed": false,
+    "scrolled": false
+   },
+   "outputs": [],
+   "source": [
+    "sess = tf.InteractiveSession()\n",
+    "init = tf.global_variables_initializer()\n",
+    "sess.run(init)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 30,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "1.21794 -0.528384 0.781685\n",
+      "1.23194 -0.554108 0.770444\n",
+      "1.24555 -0.578363 0.760364\n",
+      "1.25869 -0.601221 0.751355\n",
+      "1.27179 -0.622749 0.743265\n",
+      "1.2839 -0.643003 0.736142\n",
+      "1.29581 -0.662049 0.729792\n",
+      "1.30681 -0.679938 0.72423\n",
+      "1.31726 -0.696732 0.719333\n",
+      "1.32716 -0.712482 0.715035\n",
+      "1.3367 -0.727242 0.711253\n",
+      "1.34552 -0.741057 0.707969\n",
+      "1.35381 -0.753979 0.705108\n",
+      "1.3616 -0.766057 0.702621\n",
+      "1.36891 -0.777336 0.700463\n",
+      "1.37585 -0.787858 0.698589\n",
+      "1.38232 -0.797665 0.696972\n",
+      "1.38836 -0.806797 0.695578\n",
+      "1.39398 -0.815294 0.69438\n",
+      "1.39921 -0.823194 0.693351\n",
+      "1.40402 -0.830532 0.692471\n",
+      "1.40853 -0.837344 0.691716\n",
+      "1.41269 -0.843663 0.691072\n",
+      "1.41656 -0.849519 0.690522\n",
+      "1.42013 -0.854943 0.690054\n",
+      "1.42344 -0.859963 0.689655\n",
+      "1.42649 -0.864605 0.689316\n",
+      "1.42931 -0.868895 0.689029\n",
+      "1.43191 -0.872856 0.688785\n",
+      "1.43431 -0.876513 0.688579\n",
+      "1.43652 -0.879885 0.688405\n",
+      "1.43855 -0.882993 0.688257\n",
+      "1.44042 -0.885855 0.688133\n",
+      "1.44214 -0.88849 0.688029\n",
+      "1.44372 -0.890914 0.687941\n",
+      "1.44518 -0.893143 0.687866\n",
+      "1.44651 -0.895191 0.687804\n",
+      "1.44773 -0.897072 0.687752\n",
+      "1.44886 -0.898798 0.687708\n",
+      "1.44988 -0.900382 0.687671\n",
+      "1.45083 -0.901834 0.687641\n",
+      "1.45169 -0.903165 0.687615\n",
+      "1.45248 -0.904384 0.687593\n",
+      "1.4532 -0.9055 0.687575\n",
+      "1.45386 -0.906522 0.68756\n",
+      "1.45446 -0.907457 0.687548\n",
+      "1.45501 -0.908312 0.687537\n",
+      "1.45551 -0.909093 0.687529\n",
+      "1.45597 -0.909808 0.687522\n",
+      "1.45639 -0.91046 0.687515\n",
+      "1.45678 -0.911056 0.687511\n",
+      "1.45712 -0.9116 0.687506\n",
+      "1.45744 -0.912096 0.687503\n",
+      "1.45773 -0.912549 0.6875\n",
+      "1.458 -0.912962 0.687498\n",
+      "1.45824 -0.913338 0.687496\n",
+      "1.45846 -0.913682 0.687494\n",
+      "1.45866 -0.913995 0.687493\n",
+      "1.45884 -0.91428 0.687492\n",
+      "1.459 -0.91454 0.687491\n",
+      "1.45915 -0.914776 0.68749\n",
+      "1.45929 -0.914991 0.687489\n",
+      "1.45942 -0.915187 0.687489\n",
+      "1.45953 -0.915366 0.687488\n",
+      "1.45963 -0.915528 0.687488\n",
+      "1.45973 -0.915676 0.687488\n",
+      "1.45981 -0.91581 0.687487\n",
+      "1.45989 -0.915933 0.687487\n",
+      "1.45996 -0.916044 0.687487\n",
+      "1.46002 -0.916145 0.687487\n",
+      "1.46008 -0.916237 0.687487\n",
+      "1.46013 -0.916321 0.687487\n",
+      "1.46018 -0.916397 0.687487\n",
+      "1.46023 -0.916466 0.687487\n",
+      "1.46027 -0.916528 0.687487\n",
+      "1.4603 -0.916585 0.687487\n",
+      "1.46033 -0.916637 0.687486\n",
+      "1.46036 -0.916684 0.687486\n",
+      "1.46039 -0.916727 0.687486\n",
+      "1.46042 -0.916766 0.687486\n",
+      "1.46044 -0.916801 0.687486\n",
+      "1.46046 -0.916833 0.687486\n",
+      "1.46048 -0.916862 0.687486\n",
+      "1.46049 -0.916889 0.687486\n",
+      "1.46051 -0.916912 0.687486\n",
+      "1.46052 -0.916934 0.687486\n",
+      "1.46053 -0.916954 0.687486\n",
+      "1.46055 -0.916972 0.687486\n",
+      "1.46056 -0.916988 0.687486\n",
+      "1.46057 -0.917003 0.687486\n",
+      "1.46057 -0.917016 0.687486\n",
+      "1.46058 -0.917029 0.687486\n",
+      "1.46059 -0.91704 0.687486\n",
+      "1.46059 -0.91705 0.687486\n",
+      "1.4606 -0.917059 0.687486\n",
+      "1.46061 -0.917067 0.687486\n",
+      "1.46061 -0.917075 0.687486\n",
+      "1.46061 -0.917081 0.687486\n",
+      "1.46062 -0.917088 0.687486\n",
+      "1.46062 -0.917093 0.687486\n"
+     ]
+    }
+   ],
+   "source": [
+    "for _ in range(100):\n",
+    "    _, first_var, _, second_var = sess.run(train)\n",
+    "    loss_ = sess.run(loss)\n",
+    "    print(first_var, second_var, loss_)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.6.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

From 2f96076f2b3ebaa1652e1223279d3163cd1722e0 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Fri, 19 Jan 2018 21:17:14 -0500
Subject: [PATCH 08/42] unit test gradient update algo in notebook

---
 scrap.ipynb | 402 ++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 277 insertions(+), 125 deletions(-)

diff --git a/scrap.ipynb b/scrap.ipynb
index 2a6ddb2bc..ae49654fb 100644
--- a/scrap.ipynb
+++ b/scrap.ipynb
@@ -2,54 +2,177 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 42,
    "metadata": {
     "collapsed": true
    },
    "outputs": [],
    "source": [
     "from edward.models import Normal\n",
+    "import numpy as np\n",
     "import tensorflow as tf"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {
     "collapsed": true
    },
    "outputs": [],
    "source": [
-    "def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):\n",
-    "    t = 0.1\n",
-    "    delta = 10e-3\n",
-    "    eta = 1e-1\n",
+    "# --- manually!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 298,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "t = 0.1\n",
+    "delta = 10e-3\n",
+    "eta = 1e-1\n",
     "\n",
-    "    update_ops = []\n",
+    "vars_manual = [1., 2.]\n",
+    "grads_manual = [3.1018744, 1.5509372]\n",
+    "s_n_manual = [0., 0.]\n",
+    "n_manual = 1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 301,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [],
+   "source": [
+    "grads_manual = [2.7902498, 1.241244]\n",
+    "n_manual = 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 303,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0.07120250977985358"
+      ]
+     },
+     "execution_count": 303,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "eta * n_manual**(-.5 + delta)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 302,
+   "metadata": {
+    "collapsed": false,
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "2\n",
+      "s_n[i]: 1.6444956260437862\n",
+      "p_n_first: 0.07120250977985358\n",
+      "p_n_second: 0.438139347908\n",
+      "p_n: 0.0311966212044\n",
+      "var: 0.756364393664\n",
+      "\n",
+      "s_n[i]: 0.3705552246045456\n",
+      "p_n_first: 0.07120250977985358\n",
+      "p_n_second: 0.621607393593\n",
+      "p_n: 0.0442600065216\n",
+      "var: 1.84100417304\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(n_manual)\n",
+    "\n",
+    "for i in [0, 1]:\n",
+    "    s_n_manual[i] = (t * grads_manual[i]**2) + (1 - t)*s_n_manual[i]\n",
+    "    print('s_n[i]:', s_n_manual[i])\n",
+    "    p_n_first = (eta * n_manual**(-.5 + delta))\n",
+    "    p_n_second = (1 + np.sqrt(s_n_manual[i]))**(-1)\n",
+    "    p_n = p_n_first * p_n_second\n",
+    "    print('p_n_first:', p_n_first)\n",
+    "    print('p_n_second:', p_n_second)\n",
+    "    print('p_n:', p_n)\n",
+    "    \n",
+    "    vars_manual[i] += -p_n * grads_manual[i]\n",
+    "    print('var:', vars_manual[i])\n",
+    "    \n",
+    "    print('')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 322,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "# tensorflow..."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 347,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "t = 0.1\n",
+    "delta = 10e-3\n",
+    "eta = 1e-1\n",
+    "    \n",
+    "def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):\n",
+    "    ops = []\n",
     "    for i, (grad, var) in enumerate(grads_and_vars):\n",
-    "        update_s_n_op = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )\n",
-    "        update_ops.append(update_s_n_op)\n",
-    "        \n",
+    "        updated_s_n = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )\n",
+    "\n",
     "        p_n_first = eta * n**(-.5 + delta)\n",
-    "        p_n_second = 1 / (1 + tf.sqrt(s_n[i]))\n",
+    "        p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)\n",
     "        p_n = p_n_first * p_n_second\n",
-    "        \n",
-    "        apply_grad_op = var.assign_add(-p_n * grad)\n",
-    "        update_ops.append(apply_grad_op)\n",
-    "    return update_ops"
+    "\n",
+    "        updated_var = var.assign_add(-p_n * grad)\n",
+    "        ops.append((updated_s_n[i], p_n_first, p_n_second, p_n, updated_var))\n",
+    "#     increment_n = n.assign_add(1.)\n",
+    "#     ops.append(increment_n)\n",
+    "    return ops"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 348,
    "metadata": {
     "collapsed": false,
     "scrolled": false
    },
    "outputs": [],
    "source": [
-    "w1 = tf.Variable(tf.random_normal([]))\n",
-    "w2 = tf.Variable(tf.random_normal([]))\n",
+    "w1 = tf.Variable(tf.constant(1.))\n",
+    "w2 = tf.Variable(tf.constant(2.))\n",
     "var_list = [w1, w2]\n",
     "\n",
     "x = tf.constant([3., 4., 5.])\n",
@@ -69,7 +192,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 349,
    "metadata": {
     "collapsed": false,
     "scrolled": false
@@ -83,123 +206,152 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 350,
    "metadata": {
-    "collapsed": false
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "# starting vals:\n",
+    "\n",
+    "# grads_and_vars: [(3.1018744, 1.0), (1.5509372, 2.0)]\n",
+    "# s_n: array([ 0.,  0.], dtype=float32)\n",
+    "# n: 1.0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 351,
+   "metadata": {
+    "collapsed": false,
+    "scrolled": false
    },
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "1.21794 -0.528384 0.781685\n",
-      "1.23194 -0.554108 0.770444\n",
-      "1.24555 -0.578363 0.760364\n",
-      "1.25869 -0.601221 0.751355\n",
-      "1.27179 -0.622749 0.743265\n",
-      "1.2839 -0.643003 0.736142\n",
-      "1.29581 -0.662049 0.729792\n",
-      "1.30681 -0.679938 0.72423\n",
-      "1.31726 -0.696732 0.719333\n",
-      "1.32716 -0.712482 0.715035\n",
-      "1.3367 -0.727242 0.711253\n",
-      "1.34552 -0.741057 0.707969\n",
-      "1.35381 -0.753979 0.705108\n",
-      "1.3616 -0.766057 0.702621\n",
-      "1.36891 -0.777336 0.700463\n",
-      "1.37585 -0.787858 0.698589\n",
-      "1.38232 -0.797665 0.696972\n",
-      "1.38836 -0.806797 0.695578\n",
-      "1.39398 -0.815294 0.69438\n",
-      "1.39921 -0.823194 0.693351\n",
-      "1.40402 -0.830532 0.692471\n",
-      "1.40853 -0.837344 0.691716\n",
-      "1.41269 -0.843663 0.691072\n",
-      "1.41656 -0.849519 0.690522\n",
-      "1.42013 -0.854943 0.690054\n",
-      "1.42344 -0.859963 0.689655\n",
-      "1.42649 -0.864605 0.689316\n",
-      "1.42931 -0.868895 0.689029\n",
-      "1.43191 -0.872856 0.688785\n",
-      "1.43431 -0.876513 0.688579\n",
-      "1.43652 -0.879885 0.688405\n",
-      "1.43855 -0.882993 0.688257\n",
-      "1.44042 -0.885855 0.688133\n",
-      "1.44214 -0.88849 0.688029\n",
-      "1.44372 -0.890914 0.687941\n",
-      "1.44518 -0.893143 0.687866\n",
-      "1.44651 -0.895191 0.687804\n",
-      "1.44773 -0.897072 0.687752\n",
-      "1.44886 -0.898798 0.687708\n",
-      "1.44988 -0.900382 0.687671\n",
-      "1.45083 -0.901834 0.687641\n",
-      "1.45169 -0.903165 0.687615\n",
-      "1.45248 -0.904384 0.687593\n",
-      "1.4532 -0.9055 0.687575\n",
-      "1.45386 -0.906522 0.68756\n",
-      "1.45446 -0.907457 0.687548\n",
-      "1.45501 -0.908312 0.687537\n",
-      "1.45551 -0.909093 0.687529\n",
-      "1.45597 -0.909808 0.687522\n",
-      "1.45639 -0.91046 0.687515\n",
-      "1.45678 -0.911056 0.687511\n",
-      "1.45712 -0.9116 0.687506\n",
-      "1.45744 -0.912096 0.687503\n",
-      "1.45773 -0.912549 0.6875\n",
-      "1.458 -0.912962 0.687498\n",
-      "1.45824 -0.913338 0.687496\n",
-      "1.45846 -0.913682 0.687494\n",
-      "1.45866 -0.913995 0.687493\n",
-      "1.45884 -0.91428 0.687492\n",
-      "1.459 -0.91454 0.687491\n",
-      "1.45915 -0.914776 0.68749\n",
-      "1.45929 -0.914991 0.687489\n",
-      "1.45942 -0.915187 0.687489\n",
-      "1.45953 -0.915366 0.687488\n",
-      "1.45963 -0.915528 0.687488\n",
-      "1.45973 -0.915676 0.687488\n",
-      "1.45981 -0.91581 0.687487\n",
-      "1.45989 -0.915933 0.687487\n",
-      "1.45996 -0.916044 0.687487\n",
-      "1.46002 -0.916145 0.687487\n",
-      "1.46008 -0.916237 0.687487\n",
-      "1.46013 -0.916321 0.687487\n",
-      "1.46018 -0.916397 0.687487\n",
-      "1.46023 -0.916466 0.687487\n",
-      "1.46027 -0.916528 0.687487\n",
-      "1.4603 -0.916585 0.687487\n",
-      "1.46033 -0.916637 0.687486\n",
-      "1.46036 -0.916684 0.687486\n",
-      "1.46039 -0.916727 0.687486\n",
-      "1.46042 -0.916766 0.687486\n",
-      "1.46044 -0.916801 0.687486\n",
-      "1.46046 -0.916833 0.687486\n",
-      "1.46048 -0.916862 0.687486\n",
-      "1.46049 -0.916889 0.687486\n",
-      "1.46051 -0.916912 0.687486\n",
-      "1.46052 -0.916934 0.687486\n",
-      "1.46053 -0.916954 0.687486\n",
-      "1.46055 -0.916972 0.687486\n",
-      "1.46056 -0.916988 0.687486\n",
-      "1.46057 -0.917003 0.687486\n",
-      "1.46057 -0.917016 0.687486\n",
-      "1.46058 -0.917029 0.687486\n",
-      "1.46059 -0.91704 0.687486\n",
-      "1.46059 -0.91705 0.687486\n",
-      "1.4606 -0.917059 0.687486\n",
-      "1.46061 -0.917067 0.687486\n",
-      "1.46061 -0.917075 0.687486\n",
-      "1.46061 -0.917081 0.687486\n",
-      "1.46062 -0.917088 0.687486\n",
-      "1.46062 -0.917093 0.687486\n"
+      "the_n: 1.0\n",
+      "the_s_n: [ 0.  0.]\n",
+      "the_grads_and_vars: [(3.1018744, 1.0), (1.5509372, 2.0)]\n",
+      "\n",
+      "s_n_i: 0.962162\n",
+      "p_n_first_i: 0.1\n",
+      "p_n_second_i: 0.504821\n",
+      "p_n_i: 0.0504821\n",
+      "updated_var: 0.843411\n",
+      "\n",
+      "s_n_i: 0.240541\n",
+      "p_n_first_i: 0.1\n",
+      "p_n_second_i: 0.670939\n",
+      "p_n_i: 0.0670939\n",
+      "updated_var: 1.89594\n",
+      "\n",
+      "the_n: 2.0\n",
+      "the_s_n: [ 0.96216249  0.24054062]\n",
+      "the_grads_and_vars: [(2.7902498, 0.84341073), (1.241244, 1.8959416)]\n",
+      "\n",
+      "s_n_i: 1.6445\n",
+      "p_n_first_i: 0.0712025\n",
+      "p_n_second_i: 0.438139\n",
+      "p_n_i: 0.0311966\n",
+      "updated_var: 0.756364\n",
+      "\n",
+      "s_n_i: 0.370555\n",
+      "p_n_first_i: 0.0712025\n",
+      "p_n_second_i: 0.621607\n",
+      "p_n_i: 0.04426\n",
+      "updated_var: 1.841\n",
+      "\n"
      ]
     }
    ],
    "source": [
-    "for _ in range(100):\n",
-    "    _, first_var, _, second_var = sess.run(train)\n",
-    "    loss_ = sess.run(loss)\n",
-    "    print(first_var, second_var, loss_)"
+    "for _ in range(2):\n",
+    "    the_n, the_s_n, the_grads_and_vars = sess.run([n, s_n, grads_and_vars])\n",
+    "    print('the_n:', the_n)\n",
+    "    print('the_s_n:', the_s_n)\n",
+    "    print('the_grads_and_vars:', the_grads_and_vars)\n",
+    "    \n",
+    "    results = sess.run(train)\n",
+    "\n",
+    "    print('')\n",
+    "    for result in results:\n",
+    "        s_n_i, p_n_first_i, p_n_second_i, p_n_i, updated_var = result\n",
+    "        print('s_n_i:', s_n_i)\n",
+    "        print('p_n_first_i:', p_n_first_i)\n",
+    "        print('p_n_second_i:', p_n_second_i)\n",
+    "        print('p_n_i:', p_n_i)\n",
+    "        print('updated_var:', updated_var)\n",
+    "        print('')\n",
+    "    sess.run(increment_n)\n",
+    "    \n",
+    "\n",
+    "# \n",
+    "# np.testing.assert_almost_equal(w1_, var_manual[0], 7)\n",
+    "# np.testing.assert_almost_equal(w2_, var_manual[1], 7)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 309,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "# after 1\n",
+    "\n",
+    "# s_n[i]: 0.962162479337536\n",
+    "# p_n_first: 0.1\n",
+    "# p_n_second: 0.504821343702\n",
+    "# p_n: 0.0504821343702\n",
+    "# var: 0.84341075974\n",
+    "\n",
+    "# s_n[i]: 0.240540619834384\n",
+    "# p_n_first: 0.1\n",
+    "# p_n_second: 0.670938574622\n",
+    "# p_n: 0.0670938574622\n",
+    "# var: 1.89594164057\n",
+    "\n",
+    "s_n_i: 0.962162\n",
+    "p_n_first_i: 0.1\n",
+    "p_n_second_i: 0.504821\n",
+    "p_n_i: 0.0504821\n",
+    "updated_var: 0.843411\n",
+    "\n",
+    "s_n_i: 0.240541\n",
+    "p_n_first_i: 0.1\n",
+    "p_n_second_i: 0.670939\n",
+    "p_n_i: 0.0670939\n",
+    "updated_var: 1.89594\n",
+    "\n",
+    "# after 2\n",
+    "\n",
+    "# 2\n",
+    "# s_n[i]: 1.6444956260437862\n",
+    "# p_n_first: 0.07120250977985358\n",
+    "# p_n_second: 0.438139347908\n",
+    "# p_n: 0.0311966212044\n",
+    "# var: 0.756364393664\n",
+    "\n",
+    "# s_n[i]: 0.3705552246045456\n",
+    "# p_n_first: 0.07120250977985358\n",
+    "# p_n_second: 0.621607393593\n",
+    "# p_n: 0.0442600065216\n",
+    "# var: 1.84100417304\n",
+    "\n",
+    "s_n_i: 1.6445\n",
+    "p_n_first_i: 0.0712025\n",
+    "p_n_second_i: 0.438139\n",
+    "p_n_i: 0.0311966\n",
+    "updated_var: 0.756364\n",
+    "\n",
+    "s_n_i: 0.370555\n",
+    "p_n_first_i: 0.0712025\n",
+    "p_n_second_i: 0.621607\n",
+    "p_n_i: 0.04426\n",
+    "updated_var: 1.841"
    ]
   }
  ],

From 2c1162b9a8ddf89f6f52f9fdf7599adaed392fb0 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 12:18:27 -0500
Subject: [PATCH 09/42] unit test gradient update algo to 3 iterations

---
 scrap.ipynb | 128 ++++++++++++++++++++++++++++++++--------------------
 1 file changed, 79 insertions(+), 49 deletions(-)

diff --git a/scrap.ipynb b/scrap.ipynb
index ae49654fb..e87484cf2 100644
--- a/scrap.ipynb
+++ b/scrap.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 42,
+   "execution_count": null,
    "metadata": {
     "collapsed": true
    },
@@ -15,7 +15,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {
     "collapsed": true
    },
@@ -26,7 +26,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 298,
+   "execution_count": 3,
    "metadata": {
     "collapsed": true
    },
@@ -44,41 +44,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 301,
+   "execution_count": 14,
    "metadata": {
     "collapsed": false
    },
    "outputs": [],
    "source": [
-    "grads_manual = [2.7902498, 1.241244]\n",
-    "n_manual = 2"
+    "# grads_manual = [2.7902498, 1.241244]\n",
+    "grads_manual = [2.6070995, 1.0711095]\n",
+    "n_manual = 3"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 303,
+   "execution_count": null,
    "metadata": {
-    "collapsed": false
+    "collapsed": true
    },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "0.07120250977985358"
-      ]
-     },
-     "execution_count": 303,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "eta * n_manual**(-.5 + delta)"
-   ]
+   "outputs": [],
+   "source": []
   },
   {
    "cell_type": "code",
-   "execution_count": 302,
+   "execution_count": 15,
    "metadata": {
     "collapsed": false,
     "scrolled": true
@@ -88,18 +76,18 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "2\n",
-      "s_n[i]: 1.6444956260437862\n",
-      "p_n_first: 0.07120250977985358\n",
-      "p_n_second: 0.438139347908\n",
-      "p_n: 0.0311966212044\n",
-      "var: 0.756364393664\n",
+      "3\n",
+      "s_n[i]: 2.1597428437294326\n",
+      "p_n_first: 0.05837280797536004\n",
+      "p_n_second: 0.404922832044\n",
+      "p_n: 0.0236364827198\n",
+      "var: 0.694741731383\n",
       "\n",
-      "s_n[i]: 0.3705552246045456\n",
-      "p_n_first: 0.07120250977985358\n",
-      "p_n_second: 0.621607393593\n",
-      "p_n: 0.0442600065216\n",
-      "var: 1.84100417304\n",
+      "s_n[i]: 0.44822725824311604\n",
+      "p_n_first: 0.05837280797536004\n",
+      "p_n_second: 0.598982532688\n",
+      "p_n: 0.0349642923612\n",
+      "var: 1.80355358733\n",
       "\n"
      ]
     }
@@ -125,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 322,
+   "execution_count": null,
    "metadata": {
     "collapsed": true
    },
@@ -136,7 +124,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 347,
+   "execution_count": 33,
    "metadata": {
     "collapsed": true
    },
@@ -164,7 +152,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 348,
+   "execution_count": 34,
    "metadata": {
     "collapsed": false,
     "scrolled": false
@@ -192,7 +180,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 349,
+   "execution_count": 39,
    "metadata": {
     "collapsed": false,
     "scrolled": false
@@ -206,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 350,
+   "execution_count": 40,
    "metadata": {
     "collapsed": true
    },
@@ -221,7 +209,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 351,
+   "execution_count": 41,
    "metadata": {
     "collapsed": false,
     "scrolled": false
@@ -231,8 +219,6 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "the_n: 1.0\n",
-      "the_s_n: [ 0.  0.]\n",
       "the_grads_and_vars: [(3.1018744, 1.0), (1.5509372, 2.0)]\n",
       "\n",
       "s_n_i: 0.962162\n",
@@ -247,8 +233,6 @@
       "p_n_i: 0.0670939\n",
       "updated_var: 1.89594\n",
       "\n",
-      "the_n: 2.0\n",
-      "the_s_n: [ 0.96216249  0.24054062]\n",
       "the_grads_and_vars: [(2.7902498, 0.84341073), (1.241244, 1.8959416)]\n",
       "\n",
       "s_n_i: 1.6445\n",
@@ -269,10 +253,11 @@
    "source": [
     "for _ in range(2):\n",
     "    the_n, the_s_n, the_grads_and_vars = sess.run([n, s_n, grads_and_vars])\n",
-    "    print('the_n:', the_n)\n",
-    "    print('the_s_n:', the_s_n)\n",
+    "#     print('the_n:', the_n)\n",
+    "#     print('the_s_n:', the_s_n)\n",
     "    print('the_grads_and_vars:', the_grads_and_vars)\n",
     "    \n",
+    "#     the_grads_and_vars, results = sess.run([grads_and_vars, train])\n",
     "    results = sess.run(train)\n",
     "\n",
     "    print('')\n",
@@ -294,7 +279,29 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 309,
+   "execution_count": 38,
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[(2.7902498, 0.75636435), (1.241244, 1.8410041)]"
+      ]
+     },
+     "execution_count": 38,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "the_grads_and_vars"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
    "metadata": {
     "collapsed": true
    },
@@ -351,8 +358,31 @@
     "p_n_first_i: 0.0712025\n",
     "p_n_second_i: 0.621607\n",
     "p_n_i: 0.04426\n",
-    "updated_var: 1.841"
+    "updated_var: 1.841\n",
+    "    \n",
+    "# after 3\n",
+    "\n",
+    "# s_n[i]: 2.1597428437294326\n",
+    "# p_n_first: 0.05837280797536004\n",
+    "# p_n_second: 0.404922832044\n",
+    "# p_n: 0.0236364827198\n",
+    "# var: 0.694741731383\n",
+    "\n",
+    "# s_n[i]: 0.44822725824311604\n",
+    "# p_n_first: 0.05837280797536004\n",
+    "# p_n_second: 0.598982532688\n",
+    "# p_n: 0.0349642923612\n",
+    "# var: 1.80355358733"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

From ad25f6d0095c10351fc26eb606f2842fd5a0db60 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 12:19:03 -0500
Subject: [PATCH 10/42] `test_kucukelbir_grad` passes

---
 tests/inferences/test_klqp.py | 52 +++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 185f6aa35..6bf0c4a5a 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -178,6 +178,58 @@ def test_rejection_sampling_klqp(self):
     # self._test_multinomial_dirichlet(
     #   ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
 
+  def test_kucukelbir_grad(self):
+    expected_grads_and_vars = [
+      [(3.1018744, 1.0), (1.5509372, 2.0)],
+      [(2.7902498, 0.84341073), (1.241244, 1.8959416)],
+      [(2.6070995, 0.694741731383), (1.0711095, 1.80355358733)]
+    ]
+    t = 0.1
+    delta = 10e-3
+    eta = 1e-1
+
+    def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
+      ops = []
+      for i, (grad, var) in enumerate(grads_and_vars):
+        updated_s_n = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )
+
+        p_n_first = eta * n**(-.5 + delta)
+        p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
+        p_n = p_n_first * p_n_second
+
+        updated_var = var.assign_add(-p_n * grad)
+        ops.append((updated_s_n[i], p_n_first, p_n_second, p_n, updated_var))
+      return ops
+
+    with self.test_session() as sess:
+      w1 = tf.Variable(tf.constant(1.))
+      w2 = tf.Variable(tf.constant(2.))
+      var_list = [w1, w2]
+
+      x = tf.constant([3., 4., 5.])
+      y = tf.constant([.8, .1, .1])
+
+      pred = tf.nn.softmax(x * w1 * w2)
+      loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))
+      grads = tf.gradients(loss, var_list)
+      grads_and_vars = list(zip(grads, var_list))
+
+      s_n = tf.Variable(tf.zeros(2))
+      n = tf.Variable(tf.constant(1.))
+
+      train = alp_optimizer_apply_gradients(n, s_n, grads_and_vars)
+      increment_n = n.assign_add(1.)
+
+      init = tf.global_variables_initializer()
+      sess.run(init)
+
+      for i in range(3):
+        actual_grads_and_vars = sess.run(grads_and_vars)
+        self.assertAllClose(
+          actual_grads_and_vars, expected_grads_and_vars[i], rtol=5e-2, atol=5e-2)
+        _ = sess.run(train)
+        _ = sess.run(increment_n)
+
 
 if __name__ == '__main__':
   ed.set_seed(42)

From 7e4a9ce012cb6a653fefb58db941e5cbb7bbe285 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 12:31:55 -0500
Subject: [PATCH 11/42] correction: `test_kucukelbir_grad` passes

---
 tests/inferences/test_klqp.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 6bf0c4a5a..6ccfe17d8 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -181,8 +181,8 @@ def test_rejection_sampling_klqp(self):
   def test_kucukelbir_grad(self):
     expected_grads_and_vars = [
       [(3.1018744, 1.0), (1.5509372, 2.0)],
-      [(2.7902498, 0.84341073), (1.241244, 1.8959416)],
-      [(2.6070995, 0.694741731383), (1.0711095, 1.80355358733)]
+      [(2.7902498, 0.8434107), (1.241244, 1.8959416)],
+      [(2.6070995, 0.7563643), (1.0711095, 1.8410041)]
     ]
     t = 0.1
     delta = 10e-3
@@ -226,7 +226,7 @@ def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
       for i in range(3):
         actual_grads_and_vars = sess.run(grads_and_vars)
         self.assertAllClose(
-          actual_grads_and_vars, expected_grads_and_vars[i], rtol=5e-2, atol=5e-2)
+          actual_grads_and_vars, expected_grads_and_vars[i], atol=1e-9)
         _ = sess.run(train)
         _ = sess.run(increment_n)
 

From 8dc4f4f564c6fa4cea3bb41d9aea71d6e9fcf5e4 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 13:25:08 -0500
Subject: [PATCH 12/42] cleanup (still skeptical this test works, as it seems
 almost stochastic

---
 tests/inferences/test_klqp.py | 43 +++++++++++++++++++----------------
 1 file changed, 23 insertions(+), 20 deletions(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 6ccfe17d8..6d83afd2d 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -198,38 +198,41 @@ def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
         p_n = p_n_first * p_n_second
 
         updated_var = var.assign_add(-p_n * grad)
-        ops.append((updated_s_n[i], p_n_first, p_n_second, p_n, updated_var))
+        ops.append(updated_var)
       return ops
 
-    with self.test_session() as sess:
-      w1 = tf.Variable(tf.constant(1.))
-      w2 = tf.Variable(tf.constant(2.))
-      var_list = [w1, w2]
+    w1 = tf.Variable(tf.constant(1.))
+    w2 = tf.Variable(tf.constant(2.))
+    var_list = [w1, w2]
 
-      x = tf.constant([3., 4., 5.])
-      y = tf.constant([.8, .1, .1])
+    x = tf.constant([3., 4., 5.])
+    y = tf.constant([.8, .1, .1])
 
-      pred = tf.nn.softmax(x * w1 * w2)
-      loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))
-      grads = tf.gradients(loss, var_list)
-      grads_and_vars = list(zip(grads, var_list))
+    pred = tf.nn.softmax(x * w1 * w2)
+    loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))
+    grads = tf.gradients(loss, var_list)
+    grads_and_vars = list(zip(grads, var_list))
 
-      s_n = tf.Variable(tf.zeros(2))
-      n = tf.Variable(tf.constant(1.))
+    s_n = tf.Variable(tf.zeros(2))
+    n = tf.Variable(tf.constant(1.))
 
-      train = alp_optimizer_apply_gradients(n, s_n, grads_and_vars)
-      increment_n = n.assign_add(1.)
+    train = alp_optimizer_apply_gradients(n, s_n, grads_and_vars)
+    increment_n = n.assign_add(1.)
+    #
 
-      init = tf.global_variables_initializer()
-      sess.run(init)
+    actual_grads_and_vars = []
 
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
       for i in range(3):
-        actual_grads_and_vars = sess.run(grads_and_vars)
-        self.assertAllClose(
-          actual_grads_and_vars, expected_grads_and_vars[i], atol=1e-9)
+        actual_grads_and_vars.append(sess.run(grads_and_vars))
         _ = sess.run(train)
         _ = sess.run(increment_n)
 
+    # import pdb; pdb.set_trace()
+    self.assertAllClose(
+      actual_grads_and_vars, expected_grads_and_vars, atol=1e-9)
+
 
 if __name__ == '__main__':
   ed.set_seed(42)

From 0aae8ede2fe19ac13e4751b929e199ccbbb36028 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 13:29:19 -0500
Subject: [PATCH 13/42] move `test_kucukelbir_grad` to separate file

---
 edward/optimizers/__init__.py | 15 +++++++++
 edward/optimizers/sgd.py      | 19 +++++++++++
 tests/inferences/test_klqp.py | 55 ------------------------------
 tests/optimizers/test_sgd.py  | 63 +++++++++++++++++++++++++++++++++++
 4 files changed, 97 insertions(+), 55 deletions(-)
 create mode 100644 edward/optimizers/__init__.py
 create mode 100644 edward/optimizers/sgd.py
 create mode 100644 tests/optimizers/test_sgd.py

diff --git a/edward/optimizers/__init__.py b/edward/optimizers/__init__.py
new file mode 100644
index 000000000..499b2df4d
--- /dev/null
+++ b/edward/optimizers/__init__.py
@@ -0,0 +1,15 @@
+"""
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from edward.optimizers.sgd import *
+
+from tensorflow.python.util.all_util import remove_undocumented
+
+_allowed_symbols = [
+    'alp_optimizer_apply_gradients',
+]
+
+remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)
diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
new file mode 100644
index 000000000..68747c68c
--- /dev/null
+++ b/edward/optimizers/sgd.py
@@ -0,0 +1,19 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
+  ops = []
+  for i, (grad, var) in enumerate(grads_and_vars):
+    updated_s_n = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )
+
+    p_n_first = eta * n**(-.5 + delta)
+    p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
+    p_n = p_n_first * p_n_second
+
+    updated_var = var.assign_add(-p_n * grad)
+    ops.append(updated_var)
+  return ops
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 6d83afd2d..185f6aa35 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -178,61 +178,6 @@ def test_rejection_sampling_klqp(self):
     # self._test_multinomial_dirichlet(
     #   ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
 
-  def test_kucukelbir_grad(self):
-    expected_grads_and_vars = [
-      [(3.1018744, 1.0), (1.5509372, 2.0)],
-      [(2.7902498, 0.8434107), (1.241244, 1.8959416)],
-      [(2.6070995, 0.7563643), (1.0711095, 1.8410041)]
-    ]
-    t = 0.1
-    delta = 10e-3
-    eta = 1e-1
-
-    def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
-      ops = []
-      for i, (grad, var) in enumerate(grads_and_vars):
-        updated_s_n = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )
-
-        p_n_first = eta * n**(-.5 + delta)
-        p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
-        p_n = p_n_first * p_n_second
-
-        updated_var = var.assign_add(-p_n * grad)
-        ops.append(updated_var)
-      return ops
-
-    w1 = tf.Variable(tf.constant(1.))
-    w2 = tf.Variable(tf.constant(2.))
-    var_list = [w1, w2]
-
-    x = tf.constant([3., 4., 5.])
-    y = tf.constant([.8, .1, .1])
-
-    pred = tf.nn.softmax(x * w1 * w2)
-    loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))
-    grads = tf.gradients(loss, var_list)
-    grads_and_vars = list(zip(grads, var_list))
-
-    s_n = tf.Variable(tf.zeros(2))
-    n = tf.Variable(tf.constant(1.))
-
-    train = alp_optimizer_apply_gradients(n, s_n, grads_and_vars)
-    increment_n = n.assign_add(1.)
-    #
-
-    actual_grads_and_vars = []
-
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      for i in range(3):
-        actual_grads_and_vars.append(sess.run(grads_and_vars))
-        _ = sess.run(train)
-        _ = sess.run(increment_n)
-
-    # import pdb; pdb.set_trace()
-    self.assertAllClose(
-      actual_grads_and_vars, expected_grads_and_vars, atol=1e-9)
-
 
 if __name__ == '__main__':
   ed.set_seed(42)
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
new file mode 100644
index 000000000..e51282a37
--- /dev/null
+++ b/tests/optimizers/test_sgd.py
@@ -0,0 +1,63 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from edward.optimizers import alp_optimizer_apply_gradients
+
+
+class test_sgd_class(tf.test.TestCase):
+
+  def test_kucukelbir_grad(self):
+    expected_grads_and_vars = [
+      [(3.1018744, 1.0), (1.5509372, 2.0)],
+      [(2.7902498, 0.8434107), (1.241244, 1.8959416)],
+      [(2.6070995, 0.7563643), (1.0711095, 1.8410041)]
+    ]
+    t = 0.1
+    delta = 10e-3
+    eta = 1e-1
+
+    def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
+      ops = []
+      for i, (grad, var) in enumerate(grads_and_vars):
+        updated_s_n = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )
+
+        p_n_first = eta * n**(-.5 + delta)
+        p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
+        p_n = p_n_first * p_n_second
+
+        updated_var = var.assign_add(-p_n * grad)
+        ops.append(updated_var)
+      return ops
+
+    w1 = tf.Variable(tf.constant(1.))
+    w2 = tf.Variable(tf.constant(2.))
+    var_list = [w1, w2]
+
+    x = tf.constant([3., 4., 5.])
+    y = tf.constant([.8, .1, .1])
+
+    pred = tf.nn.softmax(x * w1 * w2)
+    loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))
+    grads = tf.gradients(loss, var_list)
+    grads_and_vars = list(zip(grads, var_list))
+
+    s_n = tf.Variable(tf.zeros(2))
+    n = tf.Variable(tf.constant(1.))
+
+    train = alp_optimizer_apply_gradients(n, s_n, grads_and_vars)
+    increment_n = n.assign_add(1.)
+
+    actual_grads_and_vars = []
+
+    with self.test_session() as sess:
+      tf.global_variables_initializer().run()
+      for i in range(3):
+        actual_grads_and_vars.append(sess.run(grads_and_vars))
+        _ = sess.run(train)
+        _ = sess.run(increment_n)
+
+    self.assertAllClose(
+      actual_grads_and_vars, expected_grads_and_vars, atol=1e-9)

From 70172fbc55c5da4f7953786334ca2f205f21cf3e Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 13:36:52 -0500
Subject: [PATCH 14/42] add `KucukelbirOptimizer`

---
 edward/optimizers/__init__.py |  2 +-
 edward/optimizers/sgd.py      | 35 +++++++++++++++++++++++------------
 tests/optimizers/test_sgd.py  | 19 ++++---------------
 3 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/edward/optimizers/__init__.py b/edward/optimizers/__init__.py
index 499b2df4d..61642e1dc 100644
--- a/edward/optimizers/__init__.py
+++ b/edward/optimizers/__init__.py
@@ -9,7 +9,7 @@
 from tensorflow.python.util.all_util import remove_undocumented
 
 _allowed_symbols = [
-    'alp_optimizer_apply_gradients',
+    'KucukelbirOptimizer',
 ]
 
 remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)
diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
index 68747c68c..245d92f36 100644
--- a/edward/optimizers/sgd.py
+++ b/edward/optimizers/sgd.py
@@ -5,15 +5,26 @@
 import tensorflow as tf
 
 
-def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
-  ops = []
-  for i, (grad, var) in enumerate(grads_and_vars):
-    updated_s_n = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )
-
-    p_n_first = eta * n**(-.5 + delta)
-    p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
-    p_n = p_n_first * p_n_second
-
-    updated_var = var.assign_add(-p_n * grad)
-    ops.append(updated_var)
-  return ops
+class KucukelbirOptimizer:
+
+  """
+  # TODO: add me
+  """
+
+  def __init__(self, t, delta, eta):
+    self.t = t
+    self.delta = delta
+    self.eta = eta
+
+  def apply_gradients(self, n, s_n, grads_and_vars):
+    ops = []
+    for i, (grad, var) in enumerate(grads_and_vars):
+      updated_s_n = s_n[i].assign( (self.t * grad**2) + (1 - self.t) * s_n[i] )
+
+      p_n_first = self.eta * n**(-.5 + self.delta)
+      p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
+      p_n = p_n_first * p_n_second
+
+      updated_var = var.assign_add(-p_n * grad)
+      ops.append(updated_var)
+    return ops
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index e51282a37..db2266db4 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -4,7 +4,7 @@
 
 import tensorflow as tf
 
-from edward.optimizers import alp_optimizer_apply_gradients
+from edward.optimizers import KucukelbirOptimizer
 
 
 class test_sgd_class(tf.test.TestCase):
@@ -15,23 +15,11 @@ def test_kucukelbir_grad(self):
       [(2.7902498, 0.8434107), (1.241244, 1.8959416)],
       [(2.6070995, 0.7563643), (1.0711095, 1.8410041)]
     ]
+
     t = 0.1
     delta = 10e-3
     eta = 1e-1
 
-    def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
-      ops = []
-      for i, (grad, var) in enumerate(grads_and_vars):
-        updated_s_n = s_n[i].assign( (t * grad**2) + (1 - t) * s_n[i] )
-
-        p_n_first = eta * n**(-.5 + delta)
-        p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
-        p_n = p_n_first * p_n_second
-
-        updated_var = var.assign_add(-p_n * grad)
-        ops.append(updated_var)
-      return ops
-
     w1 = tf.Variable(tf.constant(1.))
     w2 = tf.Variable(tf.constant(2.))
     var_list = [w1, w2]
@@ -47,7 +35,8 @@ def alp_optimizer_apply_gradients(n, s_n, grads_and_vars):
     s_n = tf.Variable(tf.zeros(2))
     n = tf.Variable(tf.constant(1.))
 
-    train = alp_optimizer_apply_gradients(n, s_n, grads_and_vars)
+    optimizer = KucukelbirOptimizer(t=t, delta=delta, eta=eta)
+    train = optimizer.apply_gradients(n, s_n, grads_and_vars)
     increment_n = n.assign_add(1.)
 
     actual_grads_and_vars = []

From 929e25cad9eb8186f86c8ec52c9446ec1f0200c4 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 13:40:07 -0500
Subject: [PATCH 15/42] pass `n`, `s_n` into `KucukelbirOptimizer` constructor

---
 edward/optimizers/sgd.py     | 10 ++++++----
 tests/optimizers/test_sgd.py | 10 ++++++++--
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
index 245d92f36..25097aeac 100644
--- a/edward/optimizers/sgd.py
+++ b/edward/optimizers/sgd.py
@@ -11,17 +11,19 @@ class KucukelbirOptimizer:
   # TODO: add me
   """
 
-  def __init__(self, t, delta, eta):
+  def __init__(self, t, delta, eta, s_n, n):
     self.t = t
     self.delta = delta
     self.eta = eta
+    self.s_n = s_n
+    self.n = n
 
-  def apply_gradients(self, n, s_n, grads_and_vars):
+  def apply_gradients(self, grads_and_vars):
     ops = []
     for i, (grad, var) in enumerate(grads_and_vars):
-      updated_s_n = s_n[i].assign( (self.t * grad**2) + (1 - self.t) * s_n[i] )
+      updated_s_n = self.s_n[i].assign( (self.t * grad**2) + (1 - self.t) * self.s_n[i] )
 
-      p_n_first = self.eta * n**(-.5 + self.delta)
+      p_n_first = self.eta * self.n**(-.5 + self.delta)
       p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
       p_n = p_n_first * p_n_second
 
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index db2266db4..9cb3f48df 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -35,8 +35,14 @@ def test_kucukelbir_grad(self):
     s_n = tf.Variable(tf.zeros(2))
     n = tf.Variable(tf.constant(1.))
 
-    optimizer = KucukelbirOptimizer(t=t, delta=delta, eta=eta)
-    train = optimizer.apply_gradients(n, s_n, grads_and_vars)
+    optimizer = KucukelbirOptimizer(
+      t=t,
+      delta=delta,
+      eta=eta,
+      s_n=s_n,
+      n=n
+    )
+    train = optimizer.apply_gradients(grads_and_vars)
     increment_n = n.assign_add(1.)
 
     actual_grads_and_vars = []

From 95d97749c3c059ff05c6cdf2cb62b5722932ff8f Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 14:01:24 -0500
Subject: [PATCH 16/42] looking forward to seeing if this passes CI. locally, i
 have no idea what is going on. tensorflow: you exhaust me

---
 edward/optimizers/sgd.py     |  3 +++
 tests/optimizers/test_sgd.py | 24 +++++++-----------------
 2 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
index 25097aeac..5eb8e416d 100644
--- a/edward/optimizers/sgd.py
+++ b/edward/optimizers/sgd.py
@@ -8,6 +8,8 @@
 class KucukelbirOptimizer:
 
   """
+  Used for RSVI (Rejection-Sampling Variational Inference).
+
   # TODO: add me
   """
 
@@ -29,4 +31,5 @@ def apply_gradients(self, grads_and_vars):
 
       updated_var = var.assign_add(-p_n * grad)
       ops.append(updated_var)
+    ops.append(self.n.assign_add(1.))
     return ops
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index 9cb3f48df..b6ebcd253 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -16,34 +16,25 @@ def test_kucukelbir_grad(self):
       [(2.6070995, 0.7563643), (1.0711095, 1.8410041)]
     ]
 
-    t = 0.1
-    delta = 10e-3
-    eta = 1e-1
-
+    x = tf.constant([3., 4., 5.])
+    y = tf.constant([.8, .1, .1])
     w1 = tf.Variable(tf.constant(1.))
     w2 = tf.Variable(tf.constant(2.))
     var_list = [w1, w2]
 
-    x = tf.constant([3., 4., 5.])
-    y = tf.constant([.8, .1, .1])
-
     pred = tf.nn.softmax(x * w1 * w2)
     loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))
     grads = tf.gradients(loss, var_list)
     grads_and_vars = list(zip(grads, var_list))
 
-    s_n = tf.Variable(tf.zeros(2))
-    n = tf.Variable(tf.constant(1.))
-
     optimizer = KucukelbirOptimizer(
-      t=t,
-      delta=delta,
-      eta=eta,
-      s_n=s_n,
-      n=n
+      t=0.1,
+      delta=10e-3,
+      eta=1e-1,
+      s_n=tf.Variable([0., 0.]),
+      n=tf.Variable(1.)
     )
     train = optimizer.apply_gradients(grads_and_vars)
-    increment_n = n.assign_add(1.)
 
     actual_grads_and_vars = []
 
@@ -52,7 +43,6 @@ def test_kucukelbir_grad(self):
       for i in range(3):
         actual_grads_and_vars.append(sess.run(grads_and_vars))
         _ = sess.run(train)
-        _ = sess.run(increment_n)
 
     self.assertAllClose(
       actual_grads_and_vars, expected_grads_and_vars, atol=1e-9)

From c21285848c0e0cb03c2fa97b6bd654a236d32cce Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 14:24:14 -0500
Subject: [PATCH 17/42] slightly more confidence

---
 edward/inferences/inference.py             |  2 +-
 edward/inferences/klqp.py                  | 17 +++++++++++++++++
 edward/inferences/variational_inference.py | 10 ++++++----
 edward/optimizers/sgd.py                   |  4 ++--
 tests/optimizers/test_sgd.py               |  1 +
 5 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/edward/inferences/inference.py b/edward/inferences/inference.py
index a7cea84d2..ac69bcac4 100644
--- a/edward/inferences/inference.py
+++ b/edward/inferences/inference.py
@@ -123,7 +123,6 @@ def run(self, variables=None, use_coordinator=True, *args, **kwargs):
         Passed into `initialize`.
     """
     self.initialize(*args, **kwargs)
-
     if variables is None:
       init = tf.global_variables_initializer()
     else:
@@ -144,6 +143,7 @@ def run(self, variables=None, use_coordinator=True, *args, **kwargs):
 
     for _ in range(self.n_iter):
       info_dict = self.update()
+      print(info_dict)
       self.print_progress(info_dict)
 
     self.finalize()
diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 881124903..08091ee27 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -637,6 +637,23 @@ def __init__(self, latent_vars=None, data=None, rejection_sampler_vars=None):
           standard normal draws. The random variables to approximate
           must be continuous.
       """
+      if isinstance(latent_vars, list):
+        with tf.variable_scope(None, default_name="posterior"):
+          latent_vars_dict = {}
+          continuous = \
+              ('01', 'nonnegative', 'simplex', 'real', 'multivariate_real')
+          for z in latent_vars:
+            if not hasattr(z, 'support') or z.support not in continuous:
+              raise AttributeError(
+                  "Random variable {} is not continuous or a random "
+                  "variable with supported continuous support.".format(z))
+            batch_event_shape = z.batch_shape.concatenate(z.event_shape)
+            loc = tf.Variable(tf.random_normal(batch_event_shape))
+            scale = tf.nn.softplus(
+                tf.Variable(tf.random_normal(batch_event_shape)))
+            latent_vars_dict[z] = Normal(loc=loc, scale=scale)
+          latent_vars = latent_vars_dict
+          del latent_vars_dict
       super(RejectionSamplingKLqp, self).__init__(latent_vars, data)
       self.rejection_sampler_vars = rejection_sampler_vars
 
diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index 633094a3f..90daa65ac 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -117,7 +117,6 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
 
     with tf.variable_scope(None, default_name="optimizer") as scope:
       if not use_prettytensor:
-        self.grads_and_vars = grads_and_vars  # debug
         self.train = optimizer.apply_gradients(grads_and_vars,
                                                global_step=global_step)
       else:
@@ -153,8 +152,11 @@ def update(self, feed_dict=None):
 
     sess = get_session()
     # TODO: delete me
-    # grads_and_vars_debug = sess.run([self.grads_and_vars], feed_dict)
-    _, t, loss = sess.run([self.train, self.increment_t, self.loss], feed_dict)
+    import pdb; pdb.set_trace()
+    # _, t, loss = sess.run([self.train, self.increment_t, self.loss], feed_dict)
+    _, t, loss, grads_and_vars_debug = sess.run([self.train, self.increment_t, self.loss, self.grads_and_vars], feed_dict)
+    import pdb; pdb.set_trace()
+
 
     if self.debug:
       sess.run(self.op_check, feed_dict)
@@ -164,7 +166,7 @@ def update(self, feed_dict=None):
         summary = sess.run(self.summarize, feed_dict)
         self.train_writer.add_summary(summary, t)
 
-    return {'t': t, 'loss': loss}
+    return {'t': t, 'loss': loss, 'grads_and_vars_debug': grads_and_vars_debug}
 
   def print_progress(self, info_dict):
     """Print progress to output.
diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
index 5eb8e416d..a27e05fa3 100644
--- a/edward/optimizers/sgd.py
+++ b/edward/optimizers/sgd.py
@@ -30,6 +30,6 @@ def apply_gradients(self, grads_and_vars):
       p_n = p_n_first * p_n_second
 
       updated_var = var.assign_add(-p_n * grad)
-      ops.append(updated_var)
-    ops.append(self.n.assign_add(1.))
+      ops.append((grad, updated_var))
+    # ops.append((tf.add(self.n, 1.), self.n))
     return ops
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index b6ebcd253..2812511b0 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -43,6 +43,7 @@ def test_kucukelbir_grad(self):
       for i in range(3):
         actual_grads_and_vars.append(sess.run(grads_and_vars))
         _ = sess.run(train)
+        _ = sess.run(optimizer.n.assign_add(1.))
 
     self.assertAllClose(
       actual_grads_and_vars, expected_grads_and_vars, atol=1e-9)

From 81637fb30936e1c208f3e245cbf6427e3d2a2832 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 14:43:56 -0500
Subject: [PATCH 18/42] set trainable=False

---
 edward/optimizers/sgd.py     | 3 +--
 tests/optimizers/test_sgd.py | 4 ++--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
index a27e05fa3..361459c9f 100644
--- a/edward/optimizers/sgd.py
+++ b/edward/optimizers/sgd.py
@@ -30,6 +30,5 @@ def apply_gradients(self, grads_and_vars):
       p_n = p_n_first * p_n_second
 
       updated_var = var.assign_add(-p_n * grad)
-      ops.append((grad, updated_var))
-    # ops.append((tf.add(self.n, 1.), self.n))
+      ops.append(updated_var)
     return ops
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index 2812511b0..b8f1c6134 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -31,8 +31,8 @@ def test_kucukelbir_grad(self):
       t=0.1,
       delta=10e-3,
       eta=1e-1,
-      s_n=tf.Variable([0., 0.]),
-      n=tf.Variable(1.)
+      s_n=tf.Variable([0., 0.], trainable=False),
+      n=tf.Variable(1., trainable=False)
     )
     train = optimizer.apply_gradients(grads_and_vars)
 

From 7aec66ce87de1a1b609f63bcaa6975316a4f6e2d Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 21:14:39 -0500
Subject: [PATCH 19/42] initialize `n` to 0

---
 edward/optimizers/sgd.py     | 1 +
 tests/optimizers/test_sgd.py | 3 +--
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
index 361459c9f..be5fe8e41 100644
--- a/edward/optimizers/sgd.py
+++ b/edward/optimizers/sgd.py
@@ -21,6 +21,7 @@ def __init__(self, t, delta, eta, s_n, n):
     self.n = n
 
   def apply_gradients(self, grads_and_vars):
+    self.n = tf.assign_add(self.n, 1.)
     ops = []
     for i, (grad, var) in enumerate(grads_and_vars):
       updated_s_n = self.s_n[i].assign( (self.t * grad**2) + (1 - self.t) * self.s_n[i] )
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index b8f1c6134..eaa677e4a 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -32,7 +32,7 @@ def test_kucukelbir_grad(self):
       delta=10e-3,
       eta=1e-1,
       s_n=tf.Variable([0., 0.], trainable=False),
-      n=tf.Variable(1., trainable=False)
+      n=tf.Variable(0., trainable=False)
     )
     train = optimizer.apply_gradients(grads_and_vars)
 
@@ -43,7 +43,6 @@ def test_kucukelbir_grad(self):
       for i in range(3):
         actual_grads_and_vars.append(sess.run(grads_and_vars))
         _ = sess.run(train)
-        _ = sess.run(optimizer.n.assign_add(1.))
 
     self.assertAllClose(
       actual_grads_and_vars, expected_grads_and_vars, atol=1e-9)

From dda7f265f3934e5ac86af3dbfc422f1d12f9f5b8 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 21:21:00 -0500
Subject: [PATCH 20/42] assert in loop

---
 tests/optimizers/test_sgd.py | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index eaa677e4a..5cb4dc166 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -36,13 +36,10 @@ def test_kucukelbir_grad(self):
     )
     train = optimizer.apply_gradients(grads_and_vars)
 
-    actual_grads_and_vars = []
-
     with self.test_session() as sess:
       tf.global_variables_initializer().run()
       for i in range(3):
-        actual_grads_and_vars.append(sess.run(grads_and_vars))
+        actual_grads_and_vars = sess.run(grads_and_vars)
+        self.assertAllClose(
+          actual_grads_and_vars, expected_grads_and_vars[i], atol=1e-9)
         _ = sess.run(train)
-
-    self.assertAllClose(
-      actual_grads_and_vars, expected_grads_and_vars, atol=1e-9)

From 2a4ccc8cda72d38dc56a7b871feb5d7cfac270d8 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 21:32:41 -0500
Subject: [PATCH 21/42] add dummy parameter `global_step` for temporary
 compatibility

---
 edward/optimizers/sgd.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
index be5fe8e41..e9875e47c 100644
--- a/edward/optimizers/sgd.py
+++ b/edward/optimizers/sgd.py
@@ -20,7 +20,7 @@ def __init__(self, t, delta, eta, s_n, n):
     self.s_n = s_n
     self.n = n
 
-  def apply_gradients(self, grads_and_vars):
+  def apply_gradients(self, grads_and_vars, global_step=None):
     self.n = tf.assign_add(self.n, 1.)
     ops = []
     for i, (grad, var) in enumerate(grads_and_vars):

From 8f69548c48a6fbffb4a87fc04f0ae38e9be776b4 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 21:33:15 -0500
Subject: [PATCH 22/42] add `KucukelbirOptimizer`

---
 edward/inferences/variational_inference.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index 90daa65ac..cba0b3657 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -9,6 +9,7 @@
 
 from edward.inferences.inference import Inference
 from edward.models import RandomVariable
+from edward.optimizers import KucukelbirOptimizer
 from edward.util import get_session, get_variables
 
 
@@ -67,6 +68,8 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
 
     self.loss, grads_and_vars = self.build_loss_and_gradients(var_list)
 
+    self.grads_and_vars = grads_and_vars
+
     if self.logging:
       tf.summary.scalar("loss", self.loss, collections=[self._summary_key])
       for grad, var in grads_and_vars:
@@ -110,6 +113,14 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
         optimizer = tf.train.FtrlOptimizer(learning_rate)
       elif optimizer == 'rmsprop':
         optimizer = tf.train.RMSPropOptimizer(learning_rate)
+      elif optimizer == 'kucukelbir':
+        optimizer = KucukelbirOptimizer(
+          t=0.1,
+          delta=10e-3,
+          eta=1e-1,
+          s_n=tf.Variable([0., 0.], trainable=False),
+          n=tf.Variable(0., trainable=False)
+        )
       else:
         raise ValueError('Optimizer class not found:', optimizer)
     elif not isinstance(optimizer, tf.train.Optimizer):
@@ -151,11 +162,9 @@ def update(self, feed_dict=None):
         feed_dict[key] = value
 
     sess = get_session()
-    # TODO: delete me
-    import pdb; pdb.set_trace()
     # _, t, loss = sess.run([self.train, self.increment_t, self.loss], feed_dict)
+    # TODO: delete me
     _, t, loss, grads_and_vars_debug = sess.run([self.train, self.increment_t, self.loss, self.grads_and_vars], feed_dict)
-    import pdb; pdb.set_trace()
 
 
     if self.debug:

From 26f8ed82fab2156979045bff159711abcfe30646 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 21:33:58 -0500
Subject: [PATCH 23/42] 2-space indent

---
 tests/inferences/test_klqp.py | 82 +++++++++++++++++------------------
 1 file changed, 40 insertions(+), 42 deletions(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 185f6aa35..8b9309d96 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -46,65 +46,63 @@ def _test_normal_normal(self, Inference, default, *args, **kwargs):
 
   def _test_poisson_gamma(self, Inference, *args, **kwargs):
       with self.test_session() as sess:
-          x_data = np.array([2, 8, 3, 6, 1], dtype=np.int32)
+        x_data = np.array([2, 8, 3, 6, 1], dtype=np.int32)
 
-          rate = Gamma(5.0, 1.0)
-          x = Poisson(rate=rate, sample_shape=5)
+        rate = Gamma(5.0, 1.0)
+        x = Poisson(rate=rate, sample_shape=5)
 
-          qalpha = tf.nn.softplus(tf.Variable(tf.random_normal([])))
-          qbeta = tf.nn.softplus(tf.Variable(tf.random_normal([])))
-          qgamma = Gamma(qalpha, qbeta)
+        qalpha = tf.nn.softplus(tf.Variable(tf.random_normal([])))
+        qbeta = tf.nn.softplus(tf.Variable(tf.random_normal([])))
+        qgamma = Gamma(qalpha, qbeta, allow_nan_stats=False)
 
-          # Gamma rejection sampler variables
-          def gamma_reparam_func(epsilon, alpha, beta):
+        # Gamma rejection sampler variables
+        def gamma_reparam_func(epsilon, alpha, beta):
 
-            def _gamma_reparam_func(alpha=alpha, beta=beta, epsilon=epsilon):
-              a = alpha - (1. / 3)
-              b = tf.sqrt(9 * alpha - 3)
-              c = 1 + (epsilon / b)
-              z = a * c**3
-              return z
-
-            def _gamma_reparam_func_alpha_lt_1(alpha=alpha, beta=beta, epsilon=epsilon):
-              z_tilde = _gamma_reparam_func(alpha=alpha + 1, beta=beta)
-              u = np.random.uniform()
-              z = u ** (1 / alpha) * z_tilde
-              return z
+          def _gamma_reparam_func(alpha=alpha, beta=beta, epsilon=epsilon):
+            a = alpha - (1. / 3)
+            b = tf.sqrt(9 * alpha - 3)
+            c = 1 + (epsilon / b)
+            z = a * c**3
+            return z
 
-            z = tf.cond(tf.less(alpha, 1.), _gamma_reparam_func_alpha_lt_1, _gamma_reparam_func)
-            z = tf.cond(tf.equal(beta, 1.), lambda: z, lambda: tf.divide(z, beta))
+          def _gamma_reparam_func_alpha_lt_1(alpha=alpha, beta=beta, epsilon=epsilon):
+            z_tilde = _gamma_reparam_func(alpha=alpha + 1, beta=beta)
+            u = np.random.uniform()
+            z = u ** (1 / alpha) * z_tilde
             return z
 
-          gamma_rejection_sampler_vars = {
-            'reparam_func': gamma_reparam_func,
-            'epsilon_likelihood': Normal(loc=0.0, scale=1.0),
-            'm': 10.
-          }
+          z = tf.cond(tf.less(alpha, 1.), _gamma_reparam_func_alpha_lt_1, _gamma_reparam_func)
+          z = tf.cond(tf.equal(beta, 1.), lambda: z, lambda: tf.divide(z, beta))
+          return z
 
-          # sum(x_data) = 20
-          # len(x_data) = 5
-          # analytic solution: Gamma(alpha=5+20, beta=1+5)
-          inference = Inference({rate: qgamma}, data={x: x_data},
-              rejection_sampler_vars={Gamma: gamma_rejection_sampler_vars})
+        gamma_rejection_sampler_vars = {
+          'reparam_func': gamma_reparam_func,
+          'epsilon_likelihood': Normal(loc=0.0, scale=1.0),
+          'm': 10.
+        }
 
-          inference.run(*args, **kwargs)
+        # sum(x_data) = 20
+        # len(x_data) = 5
+        # analytic solution: Gamma(alpha=5+20, beta=1+5)
+        inference = Inference({rate: qgamma}, data={x: x_data},
+          rejection_sampler_vars={Gamma: gamma_rejection_sampler_vars})
 
-          import pdb; pdb.set_trace()
+        inference.run(*args, **kwargs)
 
   def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
       with self.test_session() as sess:
-          x_data = tf.constant([2, 7, 1], dtype=np.float32)
+        x_data = tf.constant([2, 7, 1], dtype=np.float32)
 
-          probs = Dirichlet([1., 1., 1.])
-          x = Multinomial(total_count=10.0, probs=probs)
+        probs = Dirichlet([1., 1., 1.])
+        x = Multinomial(total_count=10.0, probs=probs)
 
-          qalpha = tf.Variable(tf.random_normal([3]))
-          qprobs = Dirichlet(qalpha)
+        qalpha = tf.Variable(tf.random_normal([3]))
+        qprobs = Dirichlet(qalpha)
 
-          # analytic solution: Dirichlet(alpha=[1+2, 1+7, 1+1])
-          inference = Inference({probs: qprobs}, data={x: x_data})
+        # analytic solution: Dirichlet(alpha=[1+2, 1+7, 1+1])
+        inference = Inference({probs: qprobs}, data={x: x_data})
 
-          inference.run(*args, **kwargs)
+        inference.run(*args, **kwargs)
 
   def _test_model_parameter(self, Inference, *args, **kwargs):
     with self.test_session() as sess:

From c7f3ea11dd496f3c55ba1b81fa2fd2f48d81387c Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 21:34:11 -0500
Subject: [PATCH 24/42] use `KucukelbirOptimizer`

---
 tests/inferences/test_klqp.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 8b9309d96..8cafe61e0 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -172,7 +172,11 @@ def test_score_rb_klqp(self):
 
   def test_rejection_sampling_klqp(self):
     self._test_poisson_gamma(
-      ed.RejectionSamplingKLqp, n_samples=1, n_iter=5000)
+      ed.RejectionSamplingKLqp,
+      n_samples=1,
+      n_iter=500,
+      optimizer='kucukelbir'
+    )
     # self._test_multinomial_dirichlet(
     #   ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
 

From 435ec01243dac8563f401589c7694ea8d6ed260b Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sat, 20 Jan 2018 21:55:14 -0500
Subject: [PATCH 25/42] cleanup

---
 tests/optimizers/test_sgd.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
index 5cb4dc166..b2f461a1b 100644
--- a/tests/optimizers/test_sgd.py
+++ b/tests/optimizers/test_sgd.py
@@ -29,7 +29,7 @@ def test_kucukelbir_grad(self):
 
     optimizer = KucukelbirOptimizer(
       t=0.1,
-      delta=10e-3,
+      delta=10e-16,
       eta=1e-1,
       s_n=tf.Variable([0., 0.], trainable=False),
       n=tf.Variable(0., trainable=False)
@@ -40,6 +40,5 @@ def test_kucukelbir_grad(self):
       tf.global_variables_initializer().run()
       for i in range(3):
         actual_grads_and_vars = sess.run(grads_and_vars)
-        self.assertAllClose(
-          actual_grads_and_vars, expected_grads_and_vars[i], atol=1e-9)
+        self.assertAllClose(actual_grads_and_vars, expected_grads_and_vars[i], atol=1e-9)
         _ = sess.run(train)

From 45b17b8bb71f29742e34b788acd2081fd4e11606 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sun, 21 Jan 2018 11:20:01 -0500
Subject: [PATCH 26/42] test `qalpha`, `qbeta` values

---
 tests/inferences/test_klqp.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 8cafe61e0..5d5340e6d 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -89,6 +89,9 @@ def _gamma_reparam_func_alpha_lt_1(alpha=alpha, beta=beta, epsilon=epsilon):
 
         inference.run(*args, **kwargs)
 
+        self.assertAllClose(qalpha.eval(), 25., atol=1e-2)
+        self.assertAllClose(qbeta.eval(), 6., atol=1e-2)
+
   def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
       with self.test_session() as sess:
         x_data = tf.constant([2, 7, 1], dtype=np.float32)
@@ -174,7 +177,7 @@ def test_rejection_sampling_klqp(self):
     self._test_poisson_gamma(
       ed.RejectionSamplingKLqp,
       n_samples=1,
-      n_iter=500,
+      n_iter=5000,
       optimizer='kucukelbir'
     )
     # self._test_multinomial_dirichlet(

From ed6e266104500f921d09d7a28d4b78dd8bbf654c Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Sun, 21 Jan 2018 13:14:29 -0500
Subject: [PATCH 27/42] delete blank line

---
 edward/inferences/variational_inference.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index cba0b3657..15dfb5dbc 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -166,7 +166,6 @@ def update(self, feed_dict=None):
     # TODO: delete me
     _, t, loss, grads_and_vars_debug = sess.run([self.train, self.increment_t, self.loss, self.grads_and_vars], feed_dict)
 
-
     if self.debug:
       sess.run(self.op_check, feed_dict)
 

From 80cee16d842f2ca244757a7ac745799bcd8ad813 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 22 Jan 2018 21:51:38 -0500
Subject: [PATCH 28/42] add `GammaRejectionSampler`

---
 edward/inferences/variational_inference.py  |  1 -
 edward/optimizers/sgd.py                    | 35 ----------------
 edward/{optimizers => samplers}/__init__.py |  4 +-
 edward/samplers/rejection.py                | 24 +++++++++++
 tests/optimizers/test_sgd.py                | 44 ---------------------
 tests/samplers/test_rejection.py            | 19 +++++++++
 6 files changed, 45 insertions(+), 82 deletions(-)
 delete mode 100644 edward/optimizers/sgd.py
 rename edward/{optimizers => samplers}/__init__.py (80%)
 create mode 100644 edward/samplers/rejection.py
 delete mode 100644 tests/optimizers/test_sgd.py
 create mode 100644 tests/samplers/test_rejection.py

diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index 15dfb5dbc..08d4c105d 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -9,7 +9,6 @@
 
 from edward.inferences.inference import Inference
 from edward.models import RandomVariable
-from edward.optimizers import KucukelbirOptimizer
 from edward.util import get_session, get_variables
 
 
diff --git a/edward/optimizers/sgd.py b/edward/optimizers/sgd.py
deleted file mode 100644
index e9875e47c..000000000
--- a/edward/optimizers/sgd.py
+++ /dev/null
@@ -1,35 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow as tf
-
-
-class KucukelbirOptimizer:
-
-  """
-  Used for RSVI (Rejection-Sampling Variational Inference).
-
-  # TODO: add me
-  """
-
-  def __init__(self, t, delta, eta, s_n, n):
-    self.t = t
-    self.delta = delta
-    self.eta = eta
-    self.s_n = s_n
-    self.n = n
-
-  def apply_gradients(self, grads_and_vars, global_step=None):
-    self.n = tf.assign_add(self.n, 1.)
-    ops = []
-    for i, (grad, var) in enumerate(grads_and_vars):
-      updated_s_n = self.s_n[i].assign( (self.t * grad**2) + (1 - self.t) * self.s_n[i] )
-
-      p_n_first = self.eta * self.n**(-.5 + self.delta)
-      p_n_second = (1 + tf.sqrt(updated_s_n[i]))**(-1)
-      p_n = p_n_first * p_n_second
-
-      updated_var = var.assign_add(-p_n * grad)
-      ops.append(updated_var)
-    return ops
diff --git a/edward/optimizers/__init__.py b/edward/samplers/__init__.py
similarity index 80%
rename from edward/optimizers/__init__.py
rename to edward/samplers/__init__.py
index 61642e1dc..4bb20a107 100644
--- a/edward/optimizers/__init__.py
+++ b/edward/samplers/__init__.py
@@ -4,12 +4,12 @@
 from __future__ import division
 from __future__ import print_function
 
-from edward.optimizers.sgd import *
+from edward.samplers.rejection import *
 
 from tensorflow.python.util.all_util import remove_undocumented
 
 _allowed_symbols = [
-    'KucukelbirOptimizer',
+    'GammaRejectionSampler',
 ]
 
 remove_undocumented(__name__, allowed_exception_list=_allowed_symbols)
diff --git a/edward/samplers/rejection.py b/edward/samplers/rejection.py
new file mode 100644
index 000000000..b4ddd4283
--- /dev/null
+++ b/edward/samplers/rejection.py
@@ -0,0 +1,24 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+
+class GammaRejectionSampler:
+
+  @staticmethod
+  def h(epsilon, alpha, beta):
+    a = alpha - (1. / 3)
+    b = tf.sqrt(9 * alpha - 3)
+    c = 1 + (epsilon / b)
+    d = a * c**3
+    return d / beta
+
+  @staticmethod
+  def h_inverse(z, alpha, beta):
+    a = alpha - (1. / 3)
+    b = tf.sqrt(9 * alpha - 3)
+    c = beta * z / a
+    d = c ** (1 / 3)
+    return b * (d - 1)
diff --git a/tests/optimizers/test_sgd.py b/tests/optimizers/test_sgd.py
deleted file mode 100644
index b2f461a1b..000000000
--- a/tests/optimizers/test_sgd.py
+++ /dev/null
@@ -1,44 +0,0 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import tensorflow as tf
-
-from edward.optimizers import KucukelbirOptimizer
-
-
-class test_sgd_class(tf.test.TestCase):
-
-  def test_kucukelbir_grad(self):
-    expected_grads_and_vars = [
-      [(3.1018744, 1.0), (1.5509372, 2.0)],
-      [(2.7902498, 0.8434107), (1.241244, 1.8959416)],
-      [(2.6070995, 0.7563643), (1.0711095, 1.8410041)]
-    ]
-
-    x = tf.constant([3., 4., 5.])
-    y = tf.constant([.8, .1, .1])
-    w1 = tf.Variable(tf.constant(1.))
-    w2 = tf.Variable(tf.constant(2.))
-    var_list = [w1, w2]
-
-    pred = tf.nn.softmax(x * w1 * w2)
-    loss = tf.reduce_mean(-tf.reduce_sum(y*tf.log(pred)))
-    grads = tf.gradients(loss, var_list)
-    grads_and_vars = list(zip(grads, var_list))
-
-    optimizer = KucukelbirOptimizer(
-      t=0.1,
-      delta=10e-16,
-      eta=1e-1,
-      s_n=tf.Variable([0., 0.], trainable=False),
-      n=tf.Variable(0., trainable=False)
-    )
-    train = optimizer.apply_gradients(grads_and_vars)
-
-    with self.test_session() as sess:
-      tf.global_variables_initializer().run()
-      for i in range(3):
-        actual_grads_and_vars = sess.run(grads_and_vars)
-        self.assertAllClose(actual_grads_and_vars, expected_grads_and_vars[i], atol=1e-9)
-        _ = sess.run(train)
diff --git a/tests/samplers/test_rejection.py b/tests/samplers/test_rejection.py
new file mode 100644
index 000000000..84ffdffbc
--- /dev/null
+++ b/tests/samplers/test_rejection.py
@@ -0,0 +1,19 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import tensorflow as tf
+
+from edward.samplers import GammaRejectionSampler
+
+
+class test_rejection_samplers_class(tf.test.TestCase):
+
+  def test_gamma_rejection_sampler(self):
+    alpha = tf.constant(4.)
+    beta = tf.constant(2.)
+    epsilon = tf.constant(.5)
+    with self.test_session() as sess:
+      z = GammaRejectionSampler.h(epsilon, alpha, beta)
+      self.assertAllClose(GammaRejectionSampler.h_inverse(z, alpha, beta).eval(),
+        epsilon.eval(), atol=1e-6)

From ef45bc33f0bf52333001c137d066ff78134cf21e Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 22 Jan 2018 22:09:07 -0500
Subject: [PATCH 29/42] add `log_prob_s` to `GammaRejectionSampler`

---
 edward/samplers/rejection.py     | 8 +++++++-
 tests/samplers/test_rejection.py | 9 +++++++--
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/edward/samplers/rejection.py b/edward/samplers/rejection.py
index b4ddd4283..876bad907 100644
--- a/edward/samplers/rejection.py
+++ b/edward/samplers/rejection.py
@@ -2,6 +2,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import math
+
 import tensorflow as tf
 
 
@@ -20,5 +22,9 @@ def h_inverse(z, alpha, beta):
     a = alpha - (1. / 3)
     b = tf.sqrt(9 * alpha - 3)
     c = beta * z / a
-    d = c ** (1 / 3)
+    d = c**(1 / 3)
     return b * (d - 1)
+
+  @staticmethod
+  def log_prob_s(epsilon):
+    return -0.5 * (tf.log(2 * math.pi) + epsilon**2)
diff --git a/tests/samplers/test_rejection.py b/tests/samplers/test_rejection.py
index 84ffdffbc..906421b02 100644
--- a/tests/samplers/test_rejection.py
+++ b/tests/samplers/test_rejection.py
@@ -14,6 +14,11 @@ def test_gamma_rejection_sampler(self):
     beta = tf.constant(2.)
     epsilon = tf.constant(.5)
     with self.test_session() as sess:
-      z = GammaRejectionSampler.h(epsilon, alpha, beta)
-      self.assertAllClose(GammaRejectionSampler.h_inverse(z, alpha, beta).eval(),
+      sampler = GammaRejectionSampler
+      z = sampler.h(epsilon, alpha, beta)
+
+      self.assertAllClose(sampler.h_inverse(z, alpha, beta).eval(),
         epsilon.eval(), atol=1e-6)
+      # np.log(scipy.stats.norm(.5))
+      self.assertAllClose(sampler.log_prob_s(epsilon).eval(),
+        -1.0439385332046727, atol=1e-6)

From b94ef733b2ae28ea0bba5c4f501840edf8b5b87e Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 22 Jan 2018 22:59:33 -0500
Subject: [PATCH 30/42] add citation to docstring

---
 edward/samplers/rejection.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/edward/samplers/rejection.py b/edward/samplers/rejection.py
index 876bad907..2e3737b92 100644
--- a/edward/samplers/rejection.py
+++ b/edward/samplers/rejection.py
@@ -9,6 +9,8 @@
 
 class GammaRejectionSampler:
 
+  # As implemented in https://github.com/blei-lab/ars-reparameterization/blob/master/gamma/demo.ipynb
+
   @staticmethod
   def h(epsilon, alpha, beta):
     a = alpha - (1. / 3)

From a136f9ddc682494820d8363ee9293a6acb77f2de Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 22 Jan 2018 23:04:43 -0500
Subject: [PATCH 31/42] add guts of RSVI, integrating w.r.t. z

---
 edward/inferences/klqp.py | 85 ++++++++++++++-------------------------
 1 file changed, 31 insertions(+), 54 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 08091ee27..17df3e77a 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -1189,9 +1189,8 @@ def build_score_rb_loss_and_gradients(inference, var_list):
 def build_rejection_sampling_loss_and_gradients(inference, var_list):
     """
     """
-    p_log_prob = [0.0] * inference.n_samples
-    q_log_prob = [0.0] * inference.n_samples
-    r_log_prob = [0.0] * inference.n_samples
+    rep = [0.0] * inference.n_samples
+    cor = [0.0] * inference.n_samples
     base_scope = tf.get_default_graph().unique_name("inference") + '/'
     for s in range(inference.n_samples):
       # Form dictionary in order to replace conditioning on prior or
@@ -1206,84 +1205,62 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
           else:
             dict_swap[x] = qx
 
+      p_log_prob = 0.
+      q_log_prob = 0.
+      r_log_prob = 0.
+
       for z, qz in six.iteritems(inference.latent_vars):
         # Copy q(z) to obtain new set of posterior samples.
         qz_copy = copy(qz, scope=scope)
-
-        # Of course, this will evaluate to `True`. We just do this as a simple first pass.
-        if 'rsvi':
-          # --- RSVI
-
-          # Get variable shortnames
-          qz_class = qz.__class__
-          epsilon_likelihood = inference.rejection_sampler_vars[qz_class]['epsilon_likelihood']
-          reparam_func = inference.rejection_sampler_vars[qz_class]['reparam_func']
-          m = inference.rejection_sampler_vars[qz_class]['m']
-          alpha = qz.parameters['concentration']
-          beta = qz.parameters['rate']
-
-          # Sample
-
-          # TODO: pass in the real `qalpha` and `qbeta`
-          # TODO: pass in the real `qz`
-          epsilon = epsilon_likelihood.value()
-          sample = reparam_func(epsilon, alpha, beta)
-          eps_prob = epsilon_likelihood.prob(epsilon)
-          qz_prob = qz.prob(sample)
-          random_uniform = tf.random_uniform([])
-
-          # We need this line. However, let's just accept for now.
-          # if random_uniform * m * eps_prob <= qz_prob:
-
-          # RSVI ---
-        else:
-          z = qz_copy.value()
-
-        dict_swap[z] = sample
-
-        q_log_prob[s] += tf.reduce_sum(
+        dict_swap[z] = qz_copy.value()
+        q_log_prob += tf.reduce_sum(
             inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
-        r_log_prob[s] += tf.reduce_sum(
-            inference.scale.get(z, 1.0) * epsilon_likelihood.log_prob(dict_swap[z]))
 
       for z in six.iterkeys(inference.latent_vars):
         z_copy = copy(z, dict_swap, scope=scope)
-        p_log_prob[s] += tf.reduce_sum(
+        p_log_prob += tf.reduce_sum(
             inference.scale.get(z, 1.0) * z_copy.log_prob(dict_swap[z]))
 
+      for z, qz in six.iteritems(inference.latent_vars):
+        rej_sampler = samplers[qz.__class__]
+        # TODO: how do we get per-qz variables in here?
+        epsilon_ = rej_sampler.h_inverse(dict_swap[z], alpha?, beta?)
+        z_ = rej_sampler.h(epsilon, alpha?, beta?)
+        r_log_prob += -tf.gradients(z_, epsilon_)
+
       for x in six.iterkeys(inference.data):
         if isinstance(x, RandomVariable):
           x_copy = copy(x, dict_swap, scope=scope)
-          p_log_prob[s] += tf.reduce_sum(
+          p_log_prob += tf.reduce_sum(
               inference.scale.get(x, 1.0) * x_copy.log_prob(dict_swap[x]))
 
-    p_log_prob = tf.reduce_mean(p_log_prob)
-    q_log_prob = tf.reduce_mean(q_log_prob)
-    r_log_prob = tf.reduce_mean(r_log_prob)
+      rep[s] = p_log_prob
+      cor[s] = tf.stop_gradient(p_log_prob) * (q_log_prob - r_log_prob)
 
+    rep = tf.reduce_mean(rep)
+    cor = tf.reduce_mean(cor)
     q_entropy = tf.reduce_sum([
         tf.reduce_sum(qz.entropy())
         for z, qz in six.iteritems(inference.latent_vars)])
-
     reg_penalty = tf.reduce_sum(tf.losses.get_regularization_losses())
 
+    loss = -(rep + q_entropy - reg_penalty)
+
     if inference.logging:
-      tf.summary.scalar("loss/p_log_prob", p_log_prob,
+      tf.summary.scalar("loss/reparam_objective", rep,
+                        collections=[inference._summary_key])
+      tf.summary.scalar("loss/correction_term", cor,
                         collections=[inference._summary_key])
       tf.summary.scalar("loss/q_entropy", q_entropy,
                         collections=[inference._summary_key])
       tf.summary.scalar("loss/reg_penalty", reg_penalty,
                         collections=[inference._summary_key])
 
-    loss = -(p_log_prob + q_entropy - reg_penalty)
-
-    # RSVI gradient components
-    model_grad = tf.gradients(p_log_prob, sample)[0]
-    q_entropy_grad = tf.gradients(q_entropy, var_list)
-    g_rep = [model_grad * grad for grad in tf.gradients(sample, var_list)]
-    g_cor = [p_log_prob * grad for grad in tf.gradients(q_log_prob - r_log_prob, var_list)]
-    grad_summands = zip(*[g_rep, g_cor, q_entropy_grad])
+    # TODO: fill in gradients
+    # g_rep =
+    # g_cor =
+    # g_entropy =
 
-    grads = [tf.reduce_sum(summand) for summand in grad_summands]
+    # grads = tf.gradients(loss, var_list)
     grads_and_vars = list(zip(grads, var_list))
     return loss, grads_and_vars

From 680894b0f4eead10ac553a7b5beb01099d36bd4d Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Tue, 23 Jan 2018 21:24:29 -0500
Subject: [PATCH 32/42] parametrize sampler with density

---
 edward/samplers/rejection.py     | 22 ++++++++++++----------
 tests/samplers/test_rejection.py | 10 +++++-----
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/edward/samplers/rejection.py b/edward/samplers/rejection.py
index 2e3737b92..d4a82471e 100644
--- a/edward/samplers/rejection.py
+++ b/edward/samplers/rejection.py
@@ -11,19 +11,21 @@ class GammaRejectionSampler:
 
   # As implemented in https://github.com/blei-lab/ars-reparameterization/blob/master/gamma/demo.ipynb
 
-  @staticmethod
-  def h(epsilon, alpha, beta):
-    a = alpha - (1. / 3)
-    b = tf.sqrt(9 * alpha - 3)
+  def __init__(self, density):
+    self.alpha = density.parameters['concentration']
+    self.beta = density.parameters['rate']
+
+  def h(self, epsilon):
+    a = self.alpha - (1. / 3)
+    b = tf.sqrt(9 * self.alpha - 3)
     c = 1 + (epsilon / b)
     d = a * c**3
-    return d / beta
+    return d / self.beta
 
-  @staticmethod
-  def h_inverse(z, alpha, beta):
-    a = alpha - (1. / 3)
-    b = tf.sqrt(9 * alpha - 3)
-    c = beta * z / a
+  def h_inverse(self, z):
+    a = self.alpha - (1. / 3)
+    b = tf.sqrt(9 * self.alpha - 3)
+    c = self.beta * z / a
     d = c**(1 / 3)
     return b * (d - 1)
 
diff --git a/tests/samplers/test_rejection.py b/tests/samplers/test_rejection.py
index 906421b02..517d8212a 100644
--- a/tests/samplers/test_rejection.py
+++ b/tests/samplers/test_rejection.py
@@ -4,20 +4,20 @@
 
 import tensorflow as tf
 
+from edward.models import Gamma
 from edward.samplers import GammaRejectionSampler
 
 
 class test_rejection_samplers_class(tf.test.TestCase):
 
   def test_gamma_rejection_sampler(self):
-    alpha = tf.constant(4.)
-    beta = tf.constant(2.)
+    gamma = Gamma(4., 2.)
     epsilon = tf.constant(.5)
     with self.test_session() as sess:
-      sampler = GammaRejectionSampler
-      z = sampler.h(epsilon, alpha, beta)
+      sampler = GammaRejectionSampler(density=gamma)
+      z = sampler.h(epsilon)
 
-      self.assertAllClose(sampler.h_inverse(z, alpha, beta).eval(),
+      self.assertAllClose(sampler.h_inverse(z).eval(),
         epsilon.eval(), atol=1e-6)
       # np.log(scipy.stats.norm(.5))
       self.assertAllClose(sampler.log_prob_s(epsilon).eval(),

From 47ba81c5d48a6b967b3691ee61bb8d8ad8361a7a Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Tue, 23 Jan 2018 21:34:44 -0500
Subject: [PATCH 33/42] pass density to rejection sampler; return gradients

---
 edward/inferences/klqp.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 17df3e77a..beb16d5f6 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -6,7 +6,8 @@
 import tensorflow as tf
 
 from edward.inferences.variational_inference import VariationalInference
-from edward.models import RandomVariable
+from edward.models import RandomVariable, Gamma
+from edward.samplers import GammaRejectionSampler
 from edward.util import copy, get_descendants
 
 try:
@@ -1189,6 +1190,10 @@ def build_score_rb_loss_and_gradients(inference, var_list):
 def build_rejection_sampling_loss_and_gradients(inference, var_list):
     """
     """
+    rej_samplers = {
+      Gamma: GammaRejectionSampler
+    }
+
     rep = [0.0] * inference.n_samples
     cor = [0.0] * inference.n_samples
     base_scope = tf.get_default_graph().unique_name("inference") + '/'
@@ -1222,11 +1227,9 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
             inference.scale.get(z, 1.0) * z_copy.log_prob(dict_swap[z]))
 
       for z, qz in six.iteritems(inference.latent_vars):
-        rej_sampler = samplers[qz.__class__]
-        # TODO: how do we get per-qz variables in here?
-        epsilon_ = rej_sampler.h_inverse(dict_swap[z], alpha?, beta?)
-        z_ = rej_sampler.h(epsilon, alpha?, beta?)
-        r_log_prob += -tf.gradients(z_, epsilon_)
+        sampler = rej_samplers[qz.__class__](density=qz)
+        epsilon = sampler.h_inverse(dict_swap[z])
+        r_log_prob += -tf.gradients(dict_swap[z], epsilon)
 
       for x in six.iterkeys(inference.data):
         if isinstance(x, RandomVariable):
@@ -1256,11 +1259,11 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
       tf.summary.scalar("loss/reg_penalty", reg_penalty,
                         collections=[inference._summary_key])
 
-    # TODO: fill in gradients
-    # g_rep =
-    # g_cor =
-    # g_entropy =
+    g_rep = tf.gradients(rep, var_list)
+    g_cor = tf.gradients(cor, var_list)
+    g_entropy = tf.gradients(q_entropy, var_list)
 
-    # grads = tf.gradients(loss, var_list)
+    grad_summands = zip(*[g_rep, g_cor, q_entropy_grad])
+    grads = [tf.reduce_sum(summand) for summand in grad_summands]
     grads_and_vars = list(zip(grads, var_list))
     return loss, grads_and_vars

From 26f0c3297b6840d1ce66c49e697c343917bd6cbe Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Tue, 23 Jan 2018 22:19:48 -0500
Subject: [PATCH 34/42] dict_swap[z] comes from rejection sampler, not `qz`

---
 edward/inferences/klqp.py                  | 11 +++++------
 edward/inferences/variational_inference.py | 10 +---------
 2 files changed, 6 insertions(+), 15 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index beb16d5f6..c88289f66 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -1217,20 +1217,19 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
       for z, qz in six.iteritems(inference.latent_vars):
         # Copy q(z) to obtain new set of posterior samples.
         qz_copy = copy(qz, scope=scope)
+        sampler = rej_samplers[qz_copy.__class__](density=qz)
         dict_swap[z] = qz_copy.value()
+        epsilon = sampler.h_inverse(tf.stop_gradient(dict_swap[z]))
+        dict_swap[z] = sampler.h(epsilon)
         q_log_prob += tf.reduce_sum(
             inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
+        r_log_prob += -tf.gradients(dict_swap[z], epsilon)[0]
 
       for z in six.iterkeys(inference.latent_vars):
         z_copy = copy(z, dict_swap, scope=scope)
         p_log_prob += tf.reduce_sum(
             inference.scale.get(z, 1.0) * z_copy.log_prob(dict_swap[z]))
 
-      for z, qz in six.iteritems(inference.latent_vars):
-        sampler = rej_samplers[qz.__class__](density=qz)
-        epsilon = sampler.h_inverse(dict_swap[z])
-        r_log_prob += -tf.gradients(dict_swap[z], epsilon)
-
       for x in six.iterkeys(inference.data):
         if isinstance(x, RandomVariable):
           x_copy = copy(x, dict_swap, scope=scope)
@@ -1263,7 +1262,7 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
     g_cor = tf.gradients(cor, var_list)
     g_entropy = tf.gradients(q_entropy, var_list)
 
-    grad_summands = zip(*[g_rep, g_cor, q_entropy_grad])
+    grad_summands = zip(*[g_rep, g_cor, g_entropy])
     grads = [tf.reduce_sum(summand) for summand in grad_summands]
     grads_and_vars = list(zip(grads, var_list))
     return loss, grads_and_vars
diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index 08d4c105d..d6494e35d 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -87,7 +87,7 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
       global_step = tf.Variable(0, trainable=False, name="global_step")
 
     if isinstance(global_step, tf.Variable):
-      starter_learning_rate = 0.1
+      starter_learning_rate = 0.01
       learning_rate = tf.train.exponential_decay(starter_learning_rate,
                                                  global_step,
                                                  100, 0.9, staircase=True)
@@ -112,14 +112,6 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
         optimizer = tf.train.FtrlOptimizer(learning_rate)
       elif optimizer == 'rmsprop':
         optimizer = tf.train.RMSPropOptimizer(learning_rate)
-      elif optimizer == 'kucukelbir':
-        optimizer = KucukelbirOptimizer(
-          t=0.1,
-          delta=10e-3,
-          eta=1e-1,
-          s_n=tf.Variable([0., 0.], trainable=False),
-          n=tf.Variable(0., trainable=False)
-        )
       else:
         raise ValueError('Optimizer class not found:', optimizer)
     elif not isinstance(optimizer, tf.train.Optimizer):

From 7b997e1cb64c893167a753f9d37a9349229e73b8 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Tue, 23 Jan 2018 22:20:12 -0500
Subject: [PATCH 35/42] delete gamma_rejection_sampler_vars

---
 tests/inferences/test_klqp.py | 40 ++++++-----------------------------
 1 file changed, 7 insertions(+), 33 deletions(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 5d5340e6d..ca22ac456 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -51,46 +51,19 @@ def _test_poisson_gamma(self, Inference, *args, **kwargs):
         rate = Gamma(5.0, 1.0)
         x = Poisson(rate=rate, sample_shape=5)
 
-        qalpha = tf.nn.softplus(tf.Variable(tf.random_normal([])))
-        qbeta = tf.nn.softplus(tf.Variable(tf.random_normal([])))
+        qalpha = tf.nn.softplus(tf.Variable(tf.random_normal([]), name='qalpha'))
+        qbeta = tf.nn.softplus(tf.Variable(tf.random_normal([]), name='qbeta'))
         qgamma = Gamma(qalpha, qbeta, allow_nan_stats=False)
 
-        # Gamma rejection sampler variables
-        def gamma_reparam_func(epsilon, alpha, beta):
-
-          def _gamma_reparam_func(alpha=alpha, beta=beta, epsilon=epsilon):
-            a = alpha - (1. / 3)
-            b = tf.sqrt(9 * alpha - 3)
-            c = 1 + (epsilon / b)
-            z = a * c**3
-            return z
-
-          def _gamma_reparam_func_alpha_lt_1(alpha=alpha, beta=beta, epsilon=epsilon):
-            z_tilde = _gamma_reparam_func(alpha=alpha + 1, beta=beta)
-            u = np.random.uniform()
-            z = u ** (1 / alpha) * z_tilde
-            return z
-
-          z = tf.cond(tf.less(alpha, 1.), _gamma_reparam_func_alpha_lt_1, _gamma_reparam_func)
-          z = tf.cond(tf.equal(beta, 1.), lambda: z, lambda: tf.divide(z, beta))
-          return z
-
-        gamma_rejection_sampler_vars = {
-          'reparam_func': gamma_reparam_func,
-          'epsilon_likelihood': Normal(loc=0.0, scale=1.0),
-          'm': 10.
-        }
-
         # sum(x_data) = 20
         # len(x_data) = 5
         # analytic solution: Gamma(alpha=5+20, beta=1+5)
-        inference = Inference({rate: qgamma}, data={x: x_data},
-          rejection_sampler_vars={Gamma: gamma_rejection_sampler_vars})
+        inference = Inference({rate: qgamma}, data={x: x_data})
 
         inference.run(*args, **kwargs)
 
-        self.assertAllClose(qalpha.eval(), 25., atol=1e-2)
-        self.assertAllClose(qbeta.eval(), 6., atol=1e-2)
+        self.assertAllClose(tf.nn.softplus(qalpha).eval(), 25., atol=1e-2)
+        self.assertAllClose(tf.nn.softplus(qbeta).eval(), 6., atol=1e-2)
 
   def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
       with self.test_session() as sess:
@@ -178,7 +151,8 @@ def test_rejection_sampling_klqp(self):
       ed.RejectionSamplingKLqp,
       n_samples=1,
       n_iter=5000,
-      optimizer='kucukelbir'
+      optimizer='rmsprop',
+      global_step=tf.Variable(0, trainable=False, name="global_step")
     )
     # self._test_multinomial_dirichlet(
     #   ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)

From 6108125555baab7cfd8b15021fd124a65be76a8f Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Tue, 23 Jan 2018 22:21:48 -0500
Subject: [PATCH 36/42] delete TODO

---
 edward/inferences/klqp.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index c88289f66..8bdf09836 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -617,7 +617,6 @@ def build_loss_and_gradients(self, var_list):
     return build_score_rb_loss_and_gradients(self, var_list)
 
 
-# TODO: you can probably make another base class that implements a `sample` method?
 class RejectionSamplingKLqp(VariationalInference):
 
     """

From 77e9a6c65248e10113cf03bbbff6f4b3a270c376 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 29 Jan 2018 21:32:51 -0500
Subject: [PATCH 37/42] WIP: _test_build_rejection_sampling_loss_and_gradients

---
 tests/inferences/test_klqp.py | 85 ++++++++++++++++++++++++++++++++---
 1 file changed, 78 insertions(+), 7 deletions(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index ca22ac456..92dadc81e 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -9,6 +9,8 @@
 from edward.models import Bernoulli, Normal, Dirichlet, Multinomial, \
   Gamma, Poisson
 
+from edward.inferences.klqp import build_rejection_sampling_loss_and_gradients
+
 
 class test_klqp_class(tf.test.TestCase):
 
@@ -80,6 +82,74 @@ def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
 
         inference.run(*args, **kwargs)
 
+  def _test_build_rejection_sampling_loss_and_gradients(self, *args, **kwargs):
+    with self.test_session() as sess:
+      def unwrap(theta):
+        alpha = np.exp(theta[0]) + 1
+        beta = np.exp(theta[1])
+        return alpha, beta
+
+      th = np.array([-0.52817175, -1.07296862])
+      alpha, beta = unwrap(th)
+      eps = np.array([0.86540763])
+      a0, b0 = np.float64(1.0), np.float64(1.0)
+      x = np.array([3, 3, 3, 3, 0], dtype=np.float64)
+
+      expected_g_reparam = np.array([-10.348131898560453, 31.81539831675293])
+      expected_g_score = np.array([0.30550423741109256, 0.0])
+      expected_g_entropy = np.array([0.28863888798339055, -1.0])
+
+      class GammaRejectionSampler:
+
+        def __init__(self, alpha, beta):
+          self.alpha, self.beta = alpha, beta
+
+        def H(self, epsilon):
+          a = self.alpha - (1. / 3)
+          b = tf.sqrt(9 * self.alpha - 3)
+          c = 1 + (epsilon / b)
+          d = a * c**3
+          return d / self.beta
+
+      MY_UNCONSTRAINED_ALPHA = tf.Variable(th[0])
+      MY_UNCONSTRAINED_BETA = tf.Variable(th[1])
+      var_list = [MY_UNCONSTRAINED_ALPHA, MY_UNCONSTRAINED_BETA]
+      MY_THETA = [MY_UNCONSTRAINED_ALPHA, MY_UNCONSTRAINED_BETA]
+      MY_ALPHA = tf.exp(MY_UNCONSTRAINED_ALPHA) + 1
+      MY_BETA = tf.exp(MY_UNCONSTRAINED_BETA)
+      MY_EPSILON = tf.Variable(eps)
+
+      sampler = GammaRejectionSampler(MY_ALPHA, MY_BETA)
+
+      dict_swap = {
+          'gamma': sampler.H(MY_EPSILON),
+          'x': tf.constant(x)
+      }
+      z_copy = Gamma(a0, b0)
+      x_copy = Poisson(dict_swap['gamma'])
+      qz_copy = Gamma(MY_ALPHA, MY_BETA)
+
+      r_log_prob = -tf.log(tf.gradients(dict_swap['gamma'], MY_EPSILON))
+
+      q_log_prob = qz_copy.log_prob(dict_swap['gamma'])
+      p_log_prob = tf.reduce_sum(z_copy.log_prob(dict_swap['gamma']))
+      p_log_prob += tf.reduce_sum(x_copy.log_prob(dict_swap['x']))
+
+      q_entropy = tf.reduce_sum([tf.reduce_sum(qz_copy.entropy())])
+
+      rep = p_log_prob
+      cor = tf.stop_gradient(p_log_prob) * (q_log_prob - r_log_prob)
+
+      g_rep = tf.gradients(rep, var_list)
+      g_cor = tf.gradients(cor, var_list)
+      g_entropy = tf.gradients(q_entropy, var_list)
+
+      tf.global_variables_initializer().run()
+
+      self.assertAllClose([g.eval() for g in g_rep], expected_g_reparam, rtol=1e-9, atol=1e-9)
+      self.assertAllClose([g.eval() for g in g_cor], expected_g_score, rtol=1e-9, atol=1e-9)
+      self.assertAllClose([g.eval() for g in g_entropy], expected_g_entropy, rtol=1e-9, atol=1e-9)
+
   def _test_model_parameter(self, Inference, *args, **kwargs):
     with self.test_session() as sess:
       x_data = np.array([0, 1, 0, 0, 0, 0, 0, 0, 0, 1])
@@ -147,13 +217,14 @@ def test_score_rb_klqp(self):
     self._test_model_parameter(ed.ScoreRBKLqp, n_iter=50)
 
   def test_rejection_sampling_klqp(self):
-    self._test_poisson_gamma(
-      ed.RejectionSamplingKLqp,
-      n_samples=1,
-      n_iter=5000,
-      optimizer='rmsprop',
-      global_step=tf.Variable(0, trainable=False, name="global_step")
-    )
+    self._test_build_rejection_sampling_loss_and_gradients()
+    # self._test_poisson_gamma(
+    #   ed.RejectionSamplingKLqp,
+    #   n_samples=1,
+    #   n_iter=5000,
+    #   optimizer='rmsprop',
+    #   global_step=tf.Variable(0, trainable=False, name="global_step")
+    # )
     # self._test_multinomial_dirichlet(
     #   ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
 

From 3846fa6710b72d687ede14392df609b97e70d4ea Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 29 Jan 2018 21:35:05 -0500
Subject: [PATCH 38/42] WIP: _test_build_rejection_sampling_loss_and_gradients

---
 tests/inferences/test_klqp.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 92dadc81e..33a89e36f 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -144,11 +144,14 @@ def H(self, epsilon):
       g_cor = tf.gradients(cor, var_list)
       g_entropy = tf.gradients(q_entropy, var_list)
 
+      grad_summands = zip(*[g_rep, g_cor, g_entropy])
+      grads = [tf.reduce_sum(summand) for summand in grad_summands]
+      grads_and_vars = list(zip(grads, var_list))
+
       tf.global_variables_initializer().run()
 
-      self.assertAllClose([g.eval() for g in g_rep], expected_g_reparam, rtol=1e-9, atol=1e-9)
-      self.assertAllClose([g.eval() for g in g_cor], expected_g_score, rtol=1e-9, atol=1e-9)
-      self.assertAllClose([g.eval() for g in g_entropy], expected_g_entropy, rtol=1e-9, atol=1e-9)
+      self.assertAllClose([g.eval() for g, v in grads_and_vars],
+        expected_g_reparam + expected_g_score + expected_g_entropy, rtol=1e-9, atol=1e-9)
 
   def _test_model_parameter(self, Inference, *args, **kwargs):
     with self.test_session() as sess:

From 23c33afe11e7a691f1347da1316ea1e7c7be0486 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 29 Jan 2018 22:12:26 -0500
Subject: [PATCH 39/42] WIP: _test_build_rejection_sampling_loss_and_gradients

---
 edward/inferences/klqp.py     | 40 ++++++++++++++--
 tests/inferences/test_klqp.py | 86 ++++++++++++++---------------------
 2 files changed, 70 insertions(+), 56 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 8bdf09836..d32ce25db 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -1186,9 +1186,36 @@ def build_score_rb_loss_and_gradients(inference, var_list):
   return loss, grads_and_vars
 
 
-def build_rejection_sampling_loss_and_gradients(inference, var_list):
+def build_rejection_sampling_loss_and_gradients(inference, var_list, epsilon=None):
     """
     """
+    # dict_swap = {
+    #     'gamma': sampler.H(MY_EPSILON),
+    #     'x': tf.constant(x)
+    # }
+    # # z_copy = Gamma(a0, b0)
+    # x_copy = Poisson(dict_swap['gamma'])
+    # qz_copy = Gamma(MY_ALPHA, MY_BETA)
+    #
+    # r_log_prob = -tf.log(tf.gradients(dict_swap['gamma'], MY_EPSILON))
+    #
+    # q_log_prob = qz_copy.log_prob(dict_swap['gamma'])
+    # p_log_prob = tf.reduce_sum(z_copy.log_prob(dict_swap['gamma']))
+    # p_log_prob += tf.reduce_sum(x_copy.log_prob(dict_swap['x']))
+    #
+    # q_entropy = tf.reduce_sum([tf.reduce_sum(qz_copy.entropy())])
+    #
+    # rep = p_log_prob
+    # cor = tf.stop_gradient(p_log_prob) * (q_log_prob - r_log_prob)
+    #
+    # g_rep = tf.gradients(rep, var_list)
+    # g_cor = tf.gradients(cor, var_list)
+    # g_entropy = tf.gradients(q_entropy, var_list)
+    #
+    # grad_summands = zip(*[g_rep, g_cor, g_entropy])
+    # grads = [tf.reduce_sum(summand) for summand in grad_summands]
+    # grads_and_vars = list(zip(grads, var_list))
+
     rej_samplers = {
       Gamma: GammaRejectionSampler
     }
@@ -1217,12 +1244,17 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list):
         # Copy q(z) to obtain new set of posterior samples.
         qz_copy = copy(qz, scope=scope)
         sampler = rej_samplers[qz_copy.__class__](density=qz)
-        dict_swap[z] = qz_copy.value()
-        epsilon = sampler.h_inverse(tf.stop_gradient(dict_swap[z]))
+
+        # dict_swap[z] = qz_copy.value()
+        # epsilon = sampler.h_inverse(tf.stop_gradient(dict_swap[z]))
+
+        if not epsilon:  # temporary
+          import sys; sys.exit()
+
         dict_swap[z] = sampler.h(epsilon)
         q_log_prob += tf.reduce_sum(
             inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
-        r_log_prob += -tf.gradients(dict_swap[z], epsilon)[0]
+        r_log_prob += -tf.log(tf.gradients(dict_swap[z], epsilon))
 
       for z in six.iterkeys(inference.latent_vars):
         z_copy = copy(z, dict_swap, scope=scope)
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 33a89e36f..e155d023d 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -84,74 +84,56 @@ def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
 
   def _test_build_rejection_sampling_loss_and_gradients(self, *args, **kwargs):
     with self.test_session() as sess:
-      def unwrap(theta):
-        alpha = np.exp(theta[0]) + 1
-        beta = np.exp(theta[1])
-        return alpha, beta
+      # def unwrap(theta):
+      #   alpha = np.exp(theta[0]) + 1
+      #   beta = np.exp(theta[1])
+      #   return alpha, beta
 
-      th = np.array([-0.52817175, -1.07296862])
-      alpha, beta = unwrap(th)
-      eps = np.array([0.86540763])
-      a0, b0 = np.float64(1.0), np.float64(1.0)
-      x = np.array([3, 3, 3, 3, 0], dtype=np.float64)
+      # th = np.array([-0.52817175, -1.07296862])
+      # alpha, beta = unwrap(th)
+      # eps = np.array([0.86540763])
+      # a0, b0 = np.float64(1.0), np.float64(1.0)
+      # x = np.array([3, 3, 3, 3, 0], dtype=np.float64)
 
       expected_g_reparam = np.array([-10.348131898560453, 31.81539831675293])
       expected_g_score = np.array([0.30550423741109256, 0.0])
       expected_g_entropy = np.array([0.28863888798339055, -1.0])
 
-      class GammaRejectionSampler:
+      # MY_UNCONSTRAINED_ALPHA = tf.Variable(th[0])
+      # MY_UNCONSTRAINED_BETA = tf.Variable(th[1])
+      # var_list = [MY_UNCONSTRAINED_ALPHA, MY_UNCONSTRAINED_BETA]
+      # MY_ALPHA = tf.exp(MY_UNCONSTRAINED_ALPHA) + 1
+      # MY_BETA = tf.exp(MY_UNCONSTRAINED_BETA)
+      # MY_EPSILON = tf.Variable(eps)
 
-        def __init__(self, alpha, beta):
-          self.alpha, self.beta = alpha, beta
+      x_data = np.array([3, 3, 3, 3, 0], dtype=np.float32)
 
-        def H(self, epsilon):
-          a = self.alpha - (1. / 3)
-          b = tf.sqrt(9 * self.alpha - 3)
-          c = 1 + (epsilon / b)
-          d = a * c**3
-          return d / self.beta
+      rate = Gamma(1.0, 1.0)
+      x = Poisson(rate=rate, sample_shape=5)
 
-      MY_UNCONSTRAINED_ALPHA = tf.Variable(th[0])
-      MY_UNCONSTRAINED_BETA = tf.Variable(th[1])
-      var_list = [MY_UNCONSTRAINED_ALPHA, MY_UNCONSTRAINED_BETA]
-      MY_THETA = [MY_UNCONSTRAINED_ALPHA, MY_UNCONSTRAINED_BETA]
-      MY_ALPHA = tf.exp(MY_UNCONSTRAINED_ALPHA) + 1
-      MY_BETA = tf.exp(MY_UNCONSTRAINED_BETA)
-      MY_EPSILON = tf.Variable(eps)
+      _qalpha = tf.Variable(-0.52817175, name='qalpha')
+      _qbeta = tf.Variable(-1.07296862, name='qbeta')
+      epsilon = tf.Variable(0.86540763, name='epsilon')
 
-      sampler = GammaRejectionSampler(MY_ALPHA, MY_BETA)
+      var_list = [_qalpha, _qbeta]
 
-      dict_swap = {
-          'gamma': sampler.H(MY_EPSILON),
-          'x': tf.constant(x)
-      }
-      z_copy = Gamma(a0, b0)
-      x_copy = Poisson(dict_swap['gamma'])
-      qz_copy = Gamma(MY_ALPHA, MY_BETA)
+      qalpha = tf.exp(_qalpha) + 1
+      qbeta = tf.exp(_qbeta)
+      qgamma = Gamma(qalpha, qbeta, allow_nan_stats=False)
 
-      r_log_prob = -tf.log(tf.gradients(dict_swap['gamma'], MY_EPSILON))
-
-      q_log_prob = qz_copy.log_prob(dict_swap['gamma'])
-      p_log_prob = tf.reduce_sum(z_copy.log_prob(dict_swap['gamma']))
-      p_log_prob += tf.reduce_sum(x_copy.log_prob(dict_swap['x']))
-
-      q_entropy = tf.reduce_sum([tf.reduce_sum(qz_copy.entropy())])
-
-      rep = p_log_prob
-      cor = tf.stop_gradient(p_log_prob) * (q_log_prob - r_log_prob)
-
-      g_rep = tf.gradients(rep, var_list)
-      g_cor = tf.gradients(cor, var_list)
-      g_entropy = tf.gradients(q_entropy, var_list)
-
-      grad_summands = zip(*[g_rep, g_cor, g_entropy])
-      grads = [tf.reduce_sum(summand) for summand in grad_summands]
-      grads_and_vars = list(zip(grads, var_list))
+      class DummyInference:
+        n_samples = 1
+        data = {x: x_data}
+        latent_vars = {rate: qgamma}
+        scale = {}
+        logging = False
 
       tf.global_variables_initializer().run()
 
+      loss, grads_and_vars = build_rejection_sampling_loss_and_gradients(DummyInference(), var_list, epsilon=epsilon)
+
       self.assertAllClose([g.eval() for g, v in grads_and_vars],
-        expected_g_reparam + expected_g_score + expected_g_entropy, rtol=1e-9, atol=1e-9)
+        expected_g_reparam + expected_g_score + expected_g_entropy, rtol=1e-6, atol=1e-6)
 
   def _test_model_parameter(self, Inference, *args, **kwargs):
     with self.test_session() as sess:

From 4c481a05e547c2b4181a5162ea5c57f339adb21a Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 29 Jan 2018 22:16:02 -0500
Subject: [PATCH 40/42] WIP: _test_build_rejection_sampling_loss_and_gradients

---
 edward/inferences/klqp.py     |  9 +++++----
 tests/inferences/test_klqp.py | 31 ++++++-------------------------
 2 files changed, 11 insertions(+), 29 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index d32ce25db..8cf300e39 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -1245,11 +1245,12 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list, epsilon=Non
         qz_copy = copy(qz, scope=scope)
         sampler = rej_samplers[qz_copy.__class__](density=qz)
 
-        # dict_swap[z] = qz_copy.value()
-        # epsilon = sampler.h_inverse(tf.stop_gradient(dict_swap[z]))
 
-        if not epsilon:  # temporary
-          import sys; sys.exit()
+        if epsilon is not None:  # temporary
+          pass
+        else:
+          dict_swap[z] = qz_copy.value()
+          epsilon = sampler.h_inverse(dict_swap[z])
 
         dict_swap[z] = sampler.h(epsilon)
         q_log_prob += tf.reduce_sum(
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index e155d023d..70e23679c 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -84,28 +84,6 @@ def _test_multinomial_dirichlet(self, Inference, *args, **kwargs):
 
   def _test_build_rejection_sampling_loss_and_gradients(self, *args, **kwargs):
     with self.test_session() as sess:
-      # def unwrap(theta):
-      #   alpha = np.exp(theta[0]) + 1
-      #   beta = np.exp(theta[1])
-      #   return alpha, beta
-
-      # th = np.array([-0.52817175, -1.07296862])
-      # alpha, beta = unwrap(th)
-      # eps = np.array([0.86540763])
-      # a0, b0 = np.float64(1.0), np.float64(1.0)
-      # x = np.array([3, 3, 3, 3, 0], dtype=np.float64)
-
-      expected_g_reparam = np.array([-10.348131898560453, 31.81539831675293])
-      expected_g_score = np.array([0.30550423741109256, 0.0])
-      expected_g_entropy = np.array([0.28863888798339055, -1.0])
-
-      # MY_UNCONSTRAINED_ALPHA = tf.Variable(th[0])
-      # MY_UNCONSTRAINED_BETA = tf.Variable(th[1])
-      # var_list = [MY_UNCONSTRAINED_ALPHA, MY_UNCONSTRAINED_BETA]
-      # MY_ALPHA = tf.exp(MY_UNCONSTRAINED_ALPHA) + 1
-      # MY_BETA = tf.exp(MY_UNCONSTRAINED_BETA)
-      # MY_EPSILON = tf.Variable(eps)
-
       x_data = np.array([3, 3, 3, 3, 0], dtype=np.float32)
 
       rate = Gamma(1.0, 1.0)
@@ -113,8 +91,6 @@ def _test_build_rejection_sampling_loss_and_gradients(self, *args, **kwargs):
 
       _qalpha = tf.Variable(-0.52817175, name='qalpha')
       _qbeta = tf.Variable(-1.07296862, name='qbeta')
-      epsilon = tf.Variable(0.86540763, name='epsilon')
-
       var_list = [_qalpha, _qbeta]
 
       qalpha = tf.exp(_qalpha) + 1
@@ -130,7 +106,12 @@ class DummyInference:
 
       tf.global_variables_initializer().run()
 
-      loss, grads_and_vars = build_rejection_sampling_loss_and_gradients(DummyInference(), var_list, epsilon=epsilon)
+      expected_g_reparam = np.array([-10.348131898560453, 31.81539831675293])
+      expected_g_score = np.array([0.30550423741109256, 0.0])
+      expected_g_entropy = np.array([0.28863888798339055, -1.0])
+
+      loss, grads_and_vars = build_rejection_sampling_loss_and_gradients(DummyInference(),
+        var_list, epsilon=tf.constant(0.86540763))
 
       self.assertAllClose([g.eval() for g, v in grads_and_vars],
         expected_g_reparam + expected_g_score + expected_g_entropy, rtol=1e-6, atol=1e-6)

From 00c9325531e2d12a15abb8db24decd841e618233 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 29 Jan 2018 22:22:29 -0500
Subject: [PATCH 41/42] WIP: _test_build_rejection_sampling_loss_and_gradients

---
 edward/inferences/klqp.py     | 28 ----------------------------
 tests/inferences/test_klqp.py | 14 +++++++-------
 2 files changed, 7 insertions(+), 35 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 8cf300e39..4f8245f5e 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -1189,33 +1189,6 @@ def build_score_rb_loss_and_gradients(inference, var_list):
 def build_rejection_sampling_loss_and_gradients(inference, var_list, epsilon=None):
     """
     """
-    # dict_swap = {
-    #     'gamma': sampler.H(MY_EPSILON),
-    #     'x': tf.constant(x)
-    # }
-    # # z_copy = Gamma(a0, b0)
-    # x_copy = Poisson(dict_swap['gamma'])
-    # qz_copy = Gamma(MY_ALPHA, MY_BETA)
-    #
-    # r_log_prob = -tf.log(tf.gradients(dict_swap['gamma'], MY_EPSILON))
-    #
-    # q_log_prob = qz_copy.log_prob(dict_swap['gamma'])
-    # p_log_prob = tf.reduce_sum(z_copy.log_prob(dict_swap['gamma']))
-    # p_log_prob += tf.reduce_sum(x_copy.log_prob(dict_swap['x']))
-    #
-    # q_entropy = tf.reduce_sum([tf.reduce_sum(qz_copy.entropy())])
-    #
-    # rep = p_log_prob
-    # cor = tf.stop_gradient(p_log_prob) * (q_log_prob - r_log_prob)
-    #
-    # g_rep = tf.gradients(rep, var_list)
-    # g_cor = tf.gradients(cor, var_list)
-    # g_entropy = tf.gradients(q_entropy, var_list)
-    #
-    # grad_summands = zip(*[g_rep, g_cor, g_entropy])
-    # grads = [tf.reduce_sum(summand) for summand in grad_summands]
-    # grads_and_vars = list(zip(grads, var_list))
-
     rej_samplers = {
       Gamma: GammaRejectionSampler
     }
@@ -1245,7 +1218,6 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list, epsilon=Non
         qz_copy = copy(qz, scope=scope)
         sampler = rej_samplers[qz_copy.__class__](density=qz)
 
-
         if epsilon is not None:  # temporary
           pass
         else:
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index 70e23679c..ddfafaa4b 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -184,13 +184,13 @@ def test_score_rb_klqp(self):
 
   def test_rejection_sampling_klqp(self):
     self._test_build_rejection_sampling_loss_and_gradients()
-    # self._test_poisson_gamma(
-    #   ed.RejectionSamplingKLqp,
-    #   n_samples=1,
-    #   n_iter=5000,
-    #   optimizer='rmsprop',
-    #   global_step=tf.Variable(0, trainable=False, name="global_step")
-    # )
+    self._test_poisson_gamma(
+      ed.RejectionSamplingKLqp,
+      n_samples=1,
+      n_iter=5000,
+      optimizer='rmsprop',
+      global_step=tf.Variable(0, trainable=False, name="global_step")
+    )
     # self._test_multinomial_dirichlet(
     #   ed.RejectionSamplingKLqp, n_samples=5, n_iter=5000)
 

From 40d38085d500d7b9b6df5a6826e1b982f8c4bcd6 Mon Sep 17 00:00:00 2001
From: William Wolf <wwolf@shopkeep.com>
Date: Mon, 29 Jan 2018 22:33:09 -0500
Subject: [PATCH 42/42] pep8

---
 edward/inferences/klqp.py                  | 11 ++++++-----
 edward/inferences/variational_inference.py |  2 +-
 tests/inferences/test_klqp.py              |  2 +-
 tests/samplers/test_rejection.py           |  4 ++--
 4 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/edward/inferences/klqp.py b/edward/inferences/klqp.py
index 4f8245f5e..42838c79b 100644
--- a/edward/inferences/klqp.py
+++ b/edward/inferences/klqp.py
@@ -1222,23 +1222,24 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list, epsilon=Non
           pass
         else:
           dict_swap[z] = qz_copy.value()
+          print('sample:', dict_swap[z])
           epsilon = sampler.h_inverse(dict_swap[z])
 
         dict_swap[z] = sampler.h(epsilon)
         q_log_prob += tf.reduce_sum(
-            inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
+          inference.scale.get(z, 1.0) * qz_copy.log_prob(dict_swap[z]))
         r_log_prob += -tf.log(tf.gradients(dict_swap[z], epsilon))
 
       for z in six.iterkeys(inference.latent_vars):
         z_copy = copy(z, dict_swap, scope=scope)
         p_log_prob += tf.reduce_sum(
-            inference.scale.get(z, 1.0) * z_copy.log_prob(dict_swap[z]))
+          inference.scale.get(z, 1.0) * z_copy.log_prob(dict_swap[z]))
 
       for x in six.iterkeys(inference.data):
         if isinstance(x, RandomVariable):
           x_copy = copy(x, dict_swap, scope=scope)
           p_log_prob += tf.reduce_sum(
-              inference.scale.get(x, 1.0) * x_copy.log_prob(dict_swap[x]))
+            inference.scale.get(x, 1.0) * x_copy.log_prob(dict_swap[x]))
 
       rep[s] = p_log_prob
       cor[s] = tf.stop_gradient(p_log_prob) * (q_log_prob - r_log_prob)
@@ -1246,8 +1247,8 @@ def build_rejection_sampling_loss_and_gradients(inference, var_list, epsilon=Non
     rep = tf.reduce_mean(rep)
     cor = tf.reduce_mean(cor)
     q_entropy = tf.reduce_sum([
-        tf.reduce_sum(qz.entropy())
-        for z, qz in six.iteritems(inference.latent_vars)])
+      tf.reduce_sum(qz.entropy())
+      for z, qz in six.iteritems(inference.latent_vars)])
     reg_penalty = tf.reduce_sum(tf.losses.get_regularization_losses())
 
     loss = -(rep + q_entropy - reg_penalty)
diff --git a/edward/inferences/variational_inference.py b/edward/inferences/variational_inference.py
index d6494e35d..fd1589782 100644
--- a/edward/inferences/variational_inference.py
+++ b/edward/inferences/variational_inference.py
@@ -87,7 +87,7 @@ def initialize(self, optimizer=None, var_list=None, use_prettytensor=False,
       global_step = tf.Variable(0, trainable=False, name="global_step")
 
     if isinstance(global_step, tf.Variable):
-      starter_learning_rate = 0.01
+      starter_learning_rate = 0.1
       learning_rate = tf.train.exponential_decay(starter_learning_rate,
                                                  global_step,
                                                  100, 0.9, staircase=True)
diff --git a/tests/inferences/test_klqp.py b/tests/inferences/test_klqp.py
index ddfafaa4b..81f23361a 100644
--- a/tests/inferences/test_klqp.py
+++ b/tests/inferences/test_klqp.py
@@ -187,7 +187,7 @@ def test_rejection_sampling_klqp(self):
     self._test_poisson_gamma(
       ed.RejectionSamplingKLqp,
       n_samples=1,
-      n_iter=5000,
+      n_iter=50,
       optimizer='rmsprop',
       global_step=tf.Variable(0, trainable=False, name="global_step")
     )
diff --git a/tests/samplers/test_rejection.py b/tests/samplers/test_rejection.py
index 517d8212a..3cc1f3eb4 100644
--- a/tests/samplers/test_rejection.py
+++ b/tests/samplers/test_rejection.py
@@ -11,9 +11,9 @@
 class test_rejection_samplers_class(tf.test.TestCase):
 
   def test_gamma_rejection_sampler(self):
-    gamma = Gamma(4., 2.)
-    epsilon = tf.constant(.5)
     with self.test_session() as sess:
+      gamma = Gamma(4., 2.)
+      epsilon = tf.constant(.5)
       sampler = GammaRejectionSampler(density=gamma)
       z = sampler.h(epsilon)