Fix dtype handling in tfpk.FeatureScaled.

PiperOrigin-RevId: 474410041

Author: emilyfertig

tensorflow_probability/python/internal/backend/numpy/numpy_math.py

@@ -964,7 +964,7 @@ def _softplus(x, name=None):  # pylint: disable=unused-argument
 
 sqrt = utils.copy_docstring(
     'tf.math.sqrt',
-    lambda x, name=None: np.sqrt(x))
+    lambda x, name=None: np.sqrt(_convert_to_tensor(x)))
 
 square = utils.copy_docstring(
     'tf.math.square',

tensorflow_probability/python/math/psd_kernels/BUILD

@@ -368,6 +368,7 @@ multi_substrate_py_test(
         # absl/testing:parameterized dep,
         # numpy dep,
         # tensorflow dep,
+        "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:test_util",
     ],
 )

tensorflow_probability/python/math/psd_kernels/feature_scaled.py

@@ -75,10 +75,12 @@ def __init__(
       raise ValueError(
           'Must specify exactly one of `scale_diag` and `inverse_scale_diag`.')
     with tf.name_scope(name):
+      dtype = util.maybe_get_common_dtype(
+          [kernel, scale_diag, inverse_scale_diag])
       self._scale_diag = tensor_util.convert_nonref_to_tensor(
-          scale_diag, name='scale_diag')
+          scale_diag, dtype=dtype, name='scale_diag')
       self._inverse_scale_diag = tensor_util.convert_nonref_to_tensor(
-          inverse_scale_diag, name='inverse_scale_diag')
+          inverse_scale_diag, dtype=dtype, name='inverse_scale_diag')
 
       def rescale_input(x, feature_ndims, example_ndims):
         """Computes `x / scale_diag`."""

tensorflow_probability/python/math/psd_kernels/feature_scaled_test.py

@@ -19,6 +19,7 @@
 import numpy as np
 import tensorflow.compat.v2 as tf
 
+from tensorflow_probability.python.internal import dtype_util
 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 feature_scaled
@@ -50,7 +51,7 @@ def testBatchShape(self):
     # Use 3 feature_ndims.
     kernel = exponentiated_quadratic.ExponentiatedQuadratic(
         amplitude, inner_length_scale, feature_ndims=3)
-    scale_diag = tf.ones([20, 1, 2, 1, 1, 1])
+    scale_diag = tf.ones([20, 1, 2, 1, 1, 1], dtype=self.dtype)
     ard_kernel = feature_scaled.FeatureScaled(kernel, scale_diag=scale_diag)
     self.assertAllEqual([20, 10, 2], ard_kernel.batch_shape)
     self.assertAllEqual(
@@ -77,6 +78,7 @@ def testKernelParametersBroadcast(self, feature_ndims, dims):
         2, 5, size=([3, 1, 2] + input_shape)).astype(self.dtype)
 
     ard_kernel = feature_scaled.FeatureScaled(kernel, scale_diag=length_scale)
+    self.assertIs(dtype_util.as_numpy_dtype(ard_kernel.dtype), self.dtype)
 
     x = np.random.uniform(-1, 1, size=input_shape).astype(self.dtype)
     y = np.random.uniform(-1, 1, size=input_shape).astype(self.dtype)