Skip to content

Commit

Permalink
fix binomial sample test on TPU
Browse files Browse the repository at this point in the history
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
  • Loading branch information
froystig authored and Google-ML-Automation committed Feb 4, 2025
1 parent 9f79799 commit 58a202a
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion tests/random_lax_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down

0 comments on commit 58a202a

Please sign in to comment.