Fix gamma_exponential psd kernel so that the covariance is

amplitude**2 * exp(-(||x - x'||**2 / (2 * lengthscale**2))**gamma)

rather than

amplitude**2 * exp(-(||x - x'||**(2*gamma) / (2 * lengthscale**2)))

PiperOrigin-RevId: 547842658

Author: ColCarroll

tensorflow_probability/python/math/psd_kernels/BUILD

@@ -189,6 +189,7 @@ multi_substrate_py_test(
     size = "small",
     srcs = ["gamma_exponential_test.py"],
     deps = [
+        ":exponentiated_quadratic",
         ":gamma_exponential",
         # absl/testing:parameterized dep,
         # numpy dep,

tensorflow_probability/python/math/psd_kernels/gamma_exponential.py

@@ -35,7 +35,7 @@ class GammaExponential(psd_kernel.AutoCompositeTensorPsdKernel):
 
     ```none
     k(x, y) = amplitude**2 * exp(
-        -||x - y||**(2 * power) / (2 * length_scale**2))
+        -(||x - y||**2 / (2 * length_scale**2))**power)
     ```
 
   where the double-bars represent vector length (ie, Euclidean, or L2 norm).
@@ -150,20 +150,21 @@ def _parameter_properties(cls, dtype):
   def _apply_with_distance(
       self, x1, x2, pairwise_square_distance, example_ndims=0):
 
-    if self.power is not None:
-      power = tf.convert_to_tensor(self.power)
-      power = util.pad_shape_with_ones(power, example_ndims)
-      pairwise_pow_distance = pairwise_square_distance ** power
-    else:
-      pairwise_pow_distance = pairwise_square_distance
-
-    exponent = -0.5 * pairwise_pow_distance
+    exponent = 0.5 * pairwise_square_distance
     inverse_length_scale = self._inverse_length_scale_parameter()
     if inverse_length_scale is not None:
       inverse_length_scale = util.pad_shape_with_ones(
           inverse_length_scale, example_ndims)
       exponent = exponent * tf.math.square(inverse_length_scale)
 
+    if self.power is not None:
+      power = tf.convert_to_tensor(self.power)
+      power = util.pad_shape_with_ones(power, example_ndims)
+      exponent = exponent ** power
+    else:
+      exponent = pairwise_square_distance
+    exponent = -1. * exponent
+
     if self.amplitude is not None:
       amplitude = tf.convert_to_tensor(self.amplitude)
       amplitude = util.pad_shape_with_ones(amplitude, example_ndims)

tensorflow_probability/python/math/psd_kernels/gamma_exponential_test.py

@@ -20,6 +20,7 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python.internal import test_util
+from tensorflow_probability.python.math.psd_kernels import exponentiated_quadratic
 from tensorflow_probability.python.math.psd_kernels import gamma_exponential
 
 
@@ -58,14 +59,24 @@ def testValuesAreCorrect(self, feature_ndims, dims):
       y = np.random.uniform(-1, 1, size=shape).astype(np.float32)
       self.assertAllClose(
           amplitude ** 2 * np.exp(
-              -np.float32(.5) * np.sum((x - y)**2)**gamma / length_scale**2),
+              -(np.sum((x - y)**2) /
+                (np.float32(2.) * length_scale**2))**gamma),
           self.evaluate(k.apply(x, y)))
 
   def testNoneShapes(self):
     k = gamma_exponential.GammaExponential(
         amplitude=np.reshape(np.arange(12.), [2, 3, 2]))
     self.assertAllEqual((2, 3, 2), k.batch_shape)
 
+  def testEqualsExponentiatedQuadratic(self):
+    np.random.seed(42)
+    ge = gamma_exponential.GammaExponential(
+        amplitude=3., length_scale=0.5, power=1., feature_ndims=0)
+    eq = exponentiated_quadratic.ExponentiatedQuadratic(
+        amplitude=3., length_scale=0.5, feature_ndims=0)
+    t1, t2 = np.random.rand(2, 10)
+    self.assertAllClose(ge.apply(t1, t2), eq.apply(t1, t2))
+
   def testShapesAreCorrect(self):
     k = gamma_exponential.GammaExponential(
         amplitude=1., length_scale=1., power=2.)