pretty rbf

Author: ferrine

pymc3/variational/test_functions.py

@@ -1,4 +1,4 @@
-from theano import tensor as tt
+from theano import theano, tensor as tt
 from .opvi import TestFunction
 
 
@@ -21,9 +21,9 @@ class Kernel(TestFunction):
 class RBF(Kernel):
     def __call__(self, X):
         XY = X.dot(X.T)
-        x2 = tt.reshape(tt.sum(tt.square(X), axis=1), (X.shape[0], 1))
+        x2 = tt.sum(X ** 2, axis=1).dimshuffle(0, 'x')
         X2e = tt.repeat(x2, X.shape[0], axis=1)
-        H = tt.sub(tt.add(X2e, X2e.T), 2 * XY)
+        H = X2e + X2e.T - 2. * XY
 
         V = tt.sort(H.flatten())
         length = V.shape[0]
@@ -34,9 +34,12 @@ def __call__(self, X):
                       # if odd vector
                       V[length // 2])
 
-        h = 0.5 * m / tt.log(X.shape[0].astype('float32') + 1.0)
+        h = .5 * m / tt.log(tt.cast(H.shape[0] + 1., theano.config.floatX))
 
+        #  RBF
         Kxy = tt.exp(-H / h / 2.0)
+
+        # Derivative
         dxkxy = -tt.dot(Kxy, X)
         sumkxy = tt.sum(Kxy, axis=1).dimshuffle(0, 'x')
         dxkxy = tt.add(dxkxy, tt.mul(X, sumkxy)) / h