Increase test tolerance for upcoming JAX PRNG change.

PiperOrigin-RevId: 580023819

Author: jburnim

tensorflow_probability/python/distributions/transformed_distribution_test.py

@@ -769,8 +769,7 @@ def testMVN(self, event_shape, shift, tril, dynamic_shape):
     num_samples = 7e3
     y = fake_mvn.sample(int(num_samples), seed=test_util.test_seed())
     x = y[0:5, ...]
-    self.assertAllMeansClose(y, expected_mean, axis=0,
-                             atol=0.1, rtol=0.1)
+    self.assertAllMeansClose(y, expected_mean, axis=0, atol=0.25)
     self.assertAllClose(expected_cov, sample_stats.covariance(y, sample_axis=0),
                         atol=0., rtol=0.1)
 

tensorflow_probability/python/distributions/two_piece_normal_test.py

@@ -369,7 +369,7 @@ def get_abs_sample_mean(skewness):
 
     err = self.compute_max_gradient_error(
         get_abs_sample_mean, [tf.constant(skewness, self.dtype)], delta=1e-1)
-    maxerr = 0.05 if self.dtype == np.float64 else 0.09
+    maxerr = 0.2
     self.assertLess(err, maxerr)
 
   @test_util.numpy_disable_gradient_test

tensorflow_probability/python/math/special_test.py

@@ -560,8 +560,8 @@ def _test_betaincinv_value(self, a_high, b_high, dtype, atol, rtol):
        "rtol": 2e-3},
       {"testcase_name": "float64",
        "dtype": np.float64,
-       "atol": 1e-12,
-       "rtol": 1e-11})
+       "atol": 3e-12,
+       "rtol": 3e-11})
   def testBetaincinvSmall(self, dtype, atol, rtol):
     self._test_betaincinv_value(
         a_high=1., b_high=1., dtype=dtype, atol=atol, rtol=rtol)