Fix type errors in Cauchy distribution test with JAX backend.

PiperOrigin-RevId: 285113109
davmre · Dec 12, 2019 · 3f240c5 · 3f240c5
1 parent 2922d86
commit 3f240c5
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/tensorflow_probability/python/distributions/cauchy_test.py b/tensorflow_probability/python/distributions/cauchy_test.py
@@ -75,7 +75,7 @@ def testParamStaticShapes(self):
 
   def testCauchyLogPDF(self):
     batch_size = 6
-    loc = tf.constant([3.0] * batch_size)
+    loc = tf.constant([3.0] * batch_size, dtype=tf.float32)
     scale = tf.constant([np.sqrt(10.0, dtype=np.float32)] * batch_size)
     x = np.array([-2.5, 2.5, 4.0, 0.0, -1.0, 2.0], dtype=np.float32)
     cauchy = tfd.Cauchy(loc=loc, scale=scale, validate_args=True)
@@ -104,7 +104,7 @@ def testCauchyLogPDF(self):
 
   def testCauchyLogPDFMultidimensional(self):
     batch_size = 6
-    loc = tf.constant([[3.0, -3.0]] * batch_size)
+    loc = tf.constant([[3.0, -3.0]] * batch_size, dtype=tf.float32)
     scale = tf.constant([
         [np.sqrt(10.0, dtype=np.float32),
          np.sqrt(15.0, dtype=np.float32)]] * batch_size)

diff --git a/tensorflow_probability/python/experimental/substrates/jax/distributions/BUILD b/tensorflow_probability/python/experimental/substrates/jax/distributions/BUILD
@@ -122,6 +122,7 @@ NO_TESTS = ("joint_distribution", "vector_exponential_linear_operator", "vector_
 TEST_FILENAMES = [filename + "_test" for filename in FILENAMES if filename not in NO_TESTS]
 
 TAP_WHITELIST = ["{}_test".format(filename) for filename in (
+    "cauchy",
     "deprecated_linalg",
     "exponential",
     "gamma_gamma",