adds clip fix for numerical instability

Author: bwengals

pymc3/gp/cov.py

@@ -126,13 +126,14 @@ def square_dist(self, X, Z):
         X = tt.mul(X, 1.0 / self.lengthscales)
         Xs = tt.sum(tt.square(X), 1)
         if Z is None:
-            return -2.0 * tt.dot(X, tt.transpose(X)) +\
-                   (tt.reshape(Xs, (-1, 1)) + tt.reshape(Xs, (1, -1)))
+            sqd = -2.0 * tt.dot(X, tt.transpose(X)) +\
+                  (tt.reshape(Xs, (-1, 1)) + tt.reshape(Xs, (1, -1)))
         else:
             Z = tt.mul(Z, 1.0 / self.lengthscales)
             Zs = tt.sum(tt.square(Z), 1)
-            return -2.0 * tt.dot(X, tt.transpose(Z)) +\
-                   (tt.reshape(Xs, (-1, 1)) + tt.reshape(Zs, (1, -1)))
+            sqd = -2.0 * tt.dot(X, tt.transpose(Z)) +\
+                  (tt.reshape(Xs, (-1, 1)) + tt.reshape(Zs, (1, -1)))
+        return tt.clip(sqd, 0.0, np.inf)
 
     def euclidean_dist(self, X, Z):
         r2 = self.square_dist(X, Z)
@@ -337,13 +338,14 @@ def square_dist(self, X, Z):
         X = tt.as_tensor_variable(X)
         Xs = tt.sum(tt.square(X), 1)
         if Z is None:
-            return -2.0 * tt.dot(X, tt.transpose(X)) +\
-                   (tt.reshape(Xs, (-1, 1)) + tt.reshape(Xs, (1, -1)))
+            sqd = -2.0 * tt.dot(X, tt.transpose(X)) +\
+                  (tt.reshape(Xs, (-1, 1)) + tt.reshape(Xs, (1, -1)))
         else:
             Z = tt.as_tensor_variable(Z)
             Zs = tt.sum(tt.square(Z), 1)
-            return -2.0 * tt.dot(X, tt.transpose(Z)) +\
-                   (tt.reshape(Xs, (-1, 1)) + tt.reshape(Zs, (1, -1)))
+            sqd = -2.0 * tt.dot(X, tt.transpose(Z)) +\
+                  (tt.reshape(Xs, (-1, 1)) + tt.reshape(Zs, (1, -1)))
+        return tt.clip(sqd, 0.0, np.inf)
 
     def __call__(self, X, Z=None):
         X, Z = self._slice(X, Z)

pymc3/tests/test_gp.py

@@ -169,6 +169,15 @@ def test_raises(self):
             gp.cov.ExpQuad(2, lengthscales, [True])
 
 
+class TestStability(unittest.TestCase):
+    def test_stable(self):
+        X = np.random.uniform(low=320., high=400., size=[2000,2])
+        with Model() as model:
+            cov = gp.cov.ExpQuad(2, 0.1)
+        dists = theano.function([], cov.square_dist(X, X))()
+        self.assertFalse(np.any(dists < 0))
+
+
 class TestExpQuad(unittest.TestCase):
     def test_1d(self):
         X = np.linspace(0,1,10)[:,None]