refactor stein (#2283)

* refactor stein

* fix typo

Author: ferrine

pymc3/variational/inference.py

@@ -574,7 +574,7 @@ def __init__(self, n_particles=100, jitter=.01, model=None, kernel=test_function
                 start=start, model=model, local_rv=local_rv, random_seed=random_seed)
         super(SVGD, self).__init__(
             KSD, histogram,
-            kernel,
+            kernel, op_kwargs=dict(temperature=temperature),
             model=model, random_seed=random_seed)
 
 

pymc3/variational/stein.py

@@ -20,13 +20,25 @@ def __init__(self, approx, kernel=rbf, input_matrix=None, temperature=1):
     @property
     @memoize
     def grad(self):
-        t = self.approx.normalizing_constant
-        Kxy, dxkxy = self.Kxy, self.dxkxy
-        dlogpdx = self.dlogp  # Normalized
         n = floatX(self.input_matrix.shape[0])
         temperature = self.temperature
-        svgd_grad = (tt.dot(Kxy, dlogpdx)/temperature + dxkxy/t) / n
-        return svgd_grad
+        svgd_grad = (self.density_part_grad / temperature +
+                     self.repulsive_part_grad)
+        return svgd_grad / n
+
+    @property
+    @memoize
+    def density_part_grad(self):
+        Kxy = self.Kxy
+        dlogpdx = self.dlogp
+        return tt.dot(Kxy, dlogpdx)
+
+    @property
+    @memoize
+    def repulsive_part_grad(self):
+        t = self.approx.normalizing_constant
+        dxkxy = self.dxkxy
+        return dxkxy/t
 
     @property
     def Kxy(self):