Enable most PSDKernel tests in numpy.

    - Avoid mutating numpy arrays.
    - Disable dynamic shape / dtype tests.

PiperOrigin-RevId: 463755123

Author: srvasude

tensorflow_probability/python/math/psd_kernels/BUILD

@@ -72,7 +72,6 @@ multi_substrate_py_test(
     name = "positive_semidefinite_kernel_test",
     size = "small",
     srcs = ["positive_semidefinite_kernel_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         ":positive_semidefinite_kernel",
         # absl/testing:parameterized dep,
@@ -126,7 +125,6 @@ multi_substrate_py_test(
     name = "exponentiated_quadratic_test",
     size = "small",
     srcs = ["exponentiated_quadratic_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         # absl/testing:parameterized dep,
         # numpy dep,
@@ -153,7 +151,6 @@ multi_substrate_py_test(
     name = "exp_sin_squared_test",
     size = "small",
     srcs = ["exp_sin_squared_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         # absl/testing:parameterized dep,
         # numpy dep,
@@ -181,7 +178,6 @@ multi_substrate_py_test(
     size = "medium",
     srcs = ["matern_test.py"],
     jax_size = "medium",
-    numpy_tags = ["notap"],
     shard_count = 4,
     deps = [
         # absl/testing:parameterized dep,
@@ -233,7 +229,6 @@ multi_substrate_py_test(
     name = "rational_quadratic_test",
     size = "small",
     srcs = ["rational_quadratic_test.py"],
-    numpy_tags = ["notap"],
     deps = [
         # absl/testing:parameterized dep,
         # numpy dep,
@@ -291,7 +286,6 @@ multi_substrate_py_test(
     size = "medium",
     srcs = ["schur_complement_test.py"],
     jax_size = "large",
-    numpy_tags = ["notap"],
     deps = [
         # absl/testing:parameterized dep,
         # numpy dep,

tensorflow_probability/python/math/psd_kernels/exp_sin_squared.py

@@ -98,20 +98,20 @@ def _apply(self, x1, x2, example_ndims=0):
       # pad the shape with self.feature_ndims number of ones.
       period = util.pad_shape_with_ones(
           period, ndims=(example_ndims + self.feature_ndims))
-      difference /= period
+      difference = difference / period
     log_kernel = util.sum_rightmost_ndims_preserving_shape(
         -2 * tf.sin(difference) ** 2, ndims=self.feature_ndims)
 
     if self.length_scale is not None:
       length_scale = tf.convert_to_tensor(self.length_scale)
       length_scale = util.pad_shape_with_ones(
           length_scale, ndims=example_ndims)
-      log_kernel /= length_scale ** 2
+      log_kernel = log_kernel / length_scale ** 2
 
     if self.amplitude is not None:
       amplitude = tf.convert_to_tensor(self.amplitude)
       amplitude = util.pad_shape_with_ones(amplitude, ndims=example_ndims)
-      log_kernel += 2. * tf.math.log(amplitude)
+      log_kernel = log_kernel + 2. * tf.math.log(amplitude)
     return tf.exp(log_kernel)
 
   @property

tensorflow_probability/python/math/psd_kernels/exp_sin_squared_test.py

@@ -27,6 +27,8 @@
 @test_util.test_all_tf_execution_regimes
 class ExpSinSquaredTest(test_util.TestCase):
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testMismatchedFloatTypesAreBad(self):
     tfp.math.psd_kernels.ExpSinSquared(1, 1)  # Should be OK (float32 fallback).
     tfp.math.psd_kernels.ExpSinSquared(1, np.float64(1))  # Should be OK.

tensorflow_probability/python/math/psd_kernels/exponentiated_quadratic_test.py

@@ -27,6 +27,8 @@
 @test_util.test_all_tf_execution_regimes
 class ExponentiatedQuadraticTest(test_util.TestCase):
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testMismatchedFloatTypesAreBad(self):
     tfp.math.psd_kernels.ExponentiatedQuadratic(
         1, 1)  # Should be OK (float32 fallback).

tensorflow_probability/python/math/psd_kernels/matern_test.py

@@ -31,6 +31,8 @@ class _MaternTestCase(test_util.TestCase):
   Subclasses must specify _kernel_type and _numpy_kernel.
   """
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testMismatchedFloatTypesAreBad(self):
     self._kernel_type(1., 1)  # Should be OK (float32 fallback).
     self._kernel_type(1., np.float64(1.))  # Should be OK.

tensorflow_probability/python/math/psd_kernels/positive_semidefinite_kernel_test.py

@@ -195,6 +195,8 @@ def testStaticBatchShape(self, params, shape):
   @parameterized.named_parameters(
       ('Dynamic-shape [2] kernel', [1., 2.], [2]),
       ('Dynamic-shape [2, 1] kernel', [[1.], [2.]], [2, 1]))
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='No dynamic shapes.')
   def testDynamicBatchShape(self, params, shape):
     tensor_params = tf1.placeholder_with_default(params, shape=None)
     k = TestKernel(tensor_params)
@@ -231,6 +233,8 @@ def testApplyOutputWithStaticShapes(self):
             y   # shape [3, 3]
         ).shape)
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='No dynamic shapes in numpy.')
   def testApplyOutputWithDynamicShapes(self):
     params_2_dynamic = tf1.placeholder_with_default([1., 2.], shape=None)
     k = TestKernel(params_2_dynamic)
@@ -437,6 +441,8 @@ def testDynamicShapesAndValuesOfProduct(self):
              for k in product_kernel[:, :1].kernels]),
         self.evaluate(product_kernel[:, :1].matrix(x, y)))
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testSumOfKernelsWithNoneDtypes(self):
     none_kernel = TestKernel()
     float32_kernel = TestKernel(np.float32(1))
@@ -451,6 +457,8 @@ def testSumOfKernelsWithNoneDtypes(self):
     with self.assertRaises(TypeError):
       _ = float32_kernel + float64_kernel
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testProductOfKernelsWithNoneDtypes(self):
     none_kernel = TestKernel()
     float32_kernel = TestKernel(np.float32(1))

tensorflow_probability/python/math/psd_kernels/rational_quadratic.py

@@ -148,7 +148,7 @@ def _apply_with_distance(
       scale_mixture_rate = tf.convert_to_tensor(self.scale_mixture_rate)
       power = util.pad_shape_with_ones(
           scale_mixture_rate, ndims=example_ndims)
-      pairwise_square_distance /= power
+      pairwise_square_distance = pairwise_square_distance / power
 
     log_result = -power * tf.math.log1p(pairwise_square_distance)
 

tensorflow_probability/python/math/psd_kernels/rational_quadratic_test.py

@@ -32,6 +32,8 @@ def _rational_quadratic(
     return (amplitude ** 2) * (1. + np.sum((x - y) ** 2) / (
         2 * scale_mixture_rate * length_scale ** 2)) ** (-scale_mixture_rate)
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testMismatchedFloatTypesAreBad(self):
     with self.assertRaises(TypeError):
       tfp.math.psd_kernels.RationalQuadratic(np.float32(1.), np.float64(1.))

tensorflow_probability/python/math/psd_kernels/schur_complement_test.py

@@ -56,6 +56,8 @@ def _broadcast_2(s1, s2):
 @test_util.test_all_tf_execution_regimes
 class SchurComplementTest(test_util.TestCase):
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testMismatchedFloatTypesAreBad(self):
     base_kernel = tfpk.ExponentiatedQuadratic(
         np.float64(5.), np.float64(.2))
@@ -255,6 +257,8 @@ def testNoneFixedInputs(self):
         self.evaluate(base_kernel.matrix(x, y)),
         self.evaluate(schur.matrix(x, y)))
 
+  @test_util.disable_test_for_backend(
+      disable_numpy=True, reason='DType mismatch not caught in numpy.')
   def testBaseKernelNoneDtype(self):
     # Test that we don't have problems when base_kernel has no explicit dtype
     # (ie, params are all None), but fixed_inputs has a different dtype than the