Automated Code Change

PiperOrigin-RevId: 906487960

Author: dsharletg

tensorflow_probability/python/distributions/BUILD

@@ -4943,7 +4943,7 @@ py_test(
         # hypothesis dep,
         # numpy dep,
         # six dep,
-        # tensorflow dep,
+        # tensorflow:tensorflow_no_contrib dep,
         "//tensorflow_probability/python/internal:hypothesis_testlib",
         "//tensorflow_probability/python/internal:tensor_util",
         "//tensorflow_probability/python/internal:test_util",

tensorflow_probability/python/distributions/half_normal_test.py

@@ -198,7 +198,7 @@ def testHalfNormalMeanAndMode(self):
     expected_mean = scale * np.sqrt(2.0) / np.sqrt(np.pi)
 
     self.assertAllEqual((3,), self.evaluate(halfnorm.mean()).shape)
-    self.assertAllEqual(expected_mean, self.evaluate(halfnorm.mean()))
+    self.assertAllClose(expected_mean, self.evaluate(halfnorm.mean()))
 
     self.assertAllEqual((3,), self.evaluate(halfnorm.mode()).shape)
     self.assertAllEqual([0., 0., 0.], self.evaluate(halfnorm.mode()))
@@ -209,16 +209,17 @@ def testHalfNormalVariance(self):
     expected_variance = scale**2.0 * (1.0 - 2.0 / np.pi)
 
     self.assertAllEqual((3,), self.evaluate(halfnorm.variance()).shape)
-    self.assertAllEqual(expected_variance, self.evaluate(halfnorm.variance()))
+    self.assertAllClose(expected_variance, self.evaluate(halfnorm.variance()))
 
   def testHalfNormalStandardDeviation(self):
     scale = np.array([7., 7., 7.])
     halfnorm = half_normal.HalfNormal(scale=scale, validate_args=True)
     expected_variance = scale**2.0 * (1.0 - 2.0 / np.pi)
 
     self.assertAllEqual((3,), halfnorm.stddev().shape)
-    self.assertAllEqual(
-        np.sqrt(expected_variance), self.evaluate(halfnorm.stddev()))
+    self.assertAllClose(
+        np.sqrt(expected_variance), self.evaluate(halfnorm.stddev())
+    )
 
   def testHalfNormalSample(self):
     scale = tf.constant(3.0)

tensorflow_probability/python/distributions/platform_compatibility_test.py

@@ -173,6 +173,7 @@
     'Logistic': 3e-6,
     'Multinomial': 2e-4,
     'NoncentralChi2': 1e-4,
+    'Normal': 5e-6,
     'OneHotCategorical': 1e-5,
     'PERT': 1e-5,
     'PowerSpherical': 2e-5,
@@ -206,16 +207,17 @@
     'InverseGamma': 5e-3,
     'JohnsonSU': 1e-2,
     'LambertWNormal': 1e-5,
-    'LKJ': .07,
+    'LKJ': 0.07,
     'Multinomial': 3e-4,
     'MultivariateNormalDiag': 5e-6,
     'MultivariateNormalFullCovariance': 1e-5,
     'MultivariateNormalTriL': 1e-5,
+    'Normal': 5e-6,
     'MultivariateNormalDiagPlusLowRankCovariance': 1e-4,
     'OneHotCategorical': 1e-3,  # TODO(b/163118820)
     'PERT': 5e-4,
     'Poisson': 3e-2,  # TODO(b/159999573)
-    'PowerSpherical': .003,
+    'PowerSpherical': 0.003,
     'RelaxedBernoulli': 3e-3,
     'RelaxedOneHotCategorical': 2e-3,  # TODO(b/163118820)
     'SigmoidBeta': 5e-4,

tensorflow_probability/python/experimental/fastgp/BUILD

@@ -145,6 +145,7 @@ py_test(
     shard_count = 3,
 #     strict_deps = False,  # DisableOnExport
     deps = [
+        ":fast_gp",
         ":fast_mtgp",
         # absl/testing:absltest dep,
         # jax dep,
@@ -293,6 +294,7 @@ py_test(
         # absl/testing:absltest dep,
         # absl/testing:parameterized dep,
         # jax dep,
+        # numpy dep,
         "//tensorflow_probability/python/internal/backend/jax",
     ],
 )

tensorflow_probability/python/experimental/fastgp/fast_mtgp_test.py

@@ -350,7 +350,7 @@ def slow_log_prob(amplitude, length_scale, noise):
     slow_value, slow_gradient = jax.value_and_grad(
         slow_log_prob, argnums=[0, 1, 2]
     )(self.dtype(1.0), self.dtype(1.0), self.dtype(1e-3))
-    np.testing.assert_allclose(value, slow_value, rtol=1e-3)
+    np.testing.assert_allclose(value, slow_value, rtol=5e-3)
     slow_d_amp, slow_d_length_scale, slow_d_noise = slow_gradient
     np.testing.assert_allclose(d_amp, slow_d_amp, rtol=2e-4)
     np.testing.assert_allclose(d_length_scale, slow_d_length_scale, rtol=2e-4)

tensorflow_probability/python/experimental/fastgp/preconditioners_test.py

@@ -102,7 +102,7 @@ def test_partial_cholesky_preconditioner(self):
     self.assertAlmostEqual(jnp.log(7.0), pcp.log_det(), places=5)
     atol = 1e-6
     if self.dtype == np.float32:
-      atol = 2e-2
+      atol = 1e-1
 
     np.testing.assert_allclose(
         jnp.array([3/7., 1/7.]),
@@ -242,7 +242,7 @@ def test_truncated_svd_split_preconditioner(self):
       ("identity", 2.0, 1.0),
       ("diagonal", 0.2, 0.2),
       ("rank_one", 0.3, 0.1),
-      ("partial_cholesky", 1e-6, 0.02),
+      ("partial_cholesky", 1e-6, 0.1),
       ("partial_lanczos", 1e-6, 0.1),
       ("truncated_svd", 0.2, 0.2),
       ("truncated_randomized_svd", 0.2, 0.2),
@@ -354,11 +354,11 @@ def test_get_preconditioner_with_kwargs(
       ("identity", 9000),
       ("diagonal", 9000),
       ("rank_one", 3100),
-      ("partial_cholesky", 5.0),
-      ("partial_lanczos", 9000.),
+      ("partial_cholesky", 10.0),
+      ("partial_lanczos", 9000.0),
       ("truncated_svd", 3100.0),
       ("truncated_randomized_svd", 3100.0),
-      ("diagonal_split", 9000.),
+      ("diagonal_split", 9000.0),
       ("rank_one_split", 3100),
       # partial_cholesky_split is basically broken until index permutations
       # are added to LowRankSplitPreconditioner.

tensorflow_probability/python/math/minimize_test.py

@@ -420,9 +420,11 @@ def loss_fn(loc, scale):
     final_x, _ = self.evaluate(
         minimize_stateless(
             loss_fn,
-            init=NormalParameters(loc=-1., scale=0.1),
+            init=NormalParameters(loc=np.float32(-1.0), scale=np.float32(0.1)),
             num_steps=250,
-            optimizer=optax.adam(0.1)))
+            optimizer=optax.adam(0.1),
+        )
+    )
     # Returned parameters should match `init` structure.
     self.assertAllClose(final_x.loc, 0., atol=0.1)
     self.assertAllClose(final_x.scale, 1., atol=0.2)