From 58a202ad18a222d13c0d3e0455322251463cee63 Mon Sep 17 00:00:00 2001 From: Roy Frostig Date: Tue, 4 Feb 2025 08:39:56 -0800 Subject: [PATCH] fix binomial sample test on TPU The recent partitionable Threefry upgrade affects binomial sampling under the RBG PRNG scheme because the implementation of `jax.random.binomial` derives internal subkeys with a call to `split`. This led a randomized test to fail by pushing its numeric closeness check just beyond its current relative tolerance. This is very likely a false failure, so we update the rtol. PiperOrigin-RevId: 723100174 --- tests/random_lax_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/random_lax_test.py b/tests/random_lax_test.py index c1b00439537e..49797b7f7c5b 100644 --- a/tests/random_lax_test.py +++ b/tests/random_lax_test.py @@ -1219,7 +1219,7 @@ def testBinomialSample(self, n, p, dtype): for samples in [uncompiled_samples, compiled_samples]: self._CheckChiSquared(samples.astype(int), pmf, pval=1e-3) self.assertAllClose(samples.mean(), n * p, rtol=0.025, check_dtypes=False) - self.assertAllClose(samples.var(), n * p * (1 - p) , rtol=0.035, + self.assertAllClose(samples.var(), n * p * (1 - p) , rtol=0.036, check_dtypes=False) def testBinomialCornerCases(self):