mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[key-reuse] fix flaky test

Browse Source
This commit is contained in:
Jake VanderPlas 2024-03-12 16:49:16 -07:00
parent 75dbd30a93
commit 2ba9b45277
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tests/key_reuse_test.py
View File

@ -608,13 +608,15 @@ class KeyReuseEagerTest(jtu.JaxTestCase):
traced_bits_msg = "In random_bits, argument 0 is already consumed."
def test_clone_eager(self):
key = jax.random.key(0)
key2 = jax.random.clone(key)
self.assertIsNot(key, key2)
# TODO(b/329326258): run this test under JIT
with jax.disable_jit():
key = jax.random.key(0)
key2 = jax.random.clone(key)
self.assertIsNot(key, key2)
_ = jax.random.uniform(key)
self.assertTrue(key._consumed)
self.assertFalse(key2._consumed)
_ = jax.random.uniform(key)
self.assertTrue(key._consumed)
self.assertFalse(key2._consumed)
def test_simple_reuse_nojit(self):
key = jax.random.key(0)