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

fix pjit_test:testWithCustomPRNGKey

Browse Source
This commit is contained in:
Jake VanderPlas 2023-04-25 10:52:15 -07:00
parent 6a23ae92ca
commit 3108f05eee
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tests/pjit_test.py
View File

@ -851,7 +851,7 @@ class PJitTest(jtu.BufferDonationTestCase):
def testWithCustomPRNGKey(self):
if not config.jax_enable_custom_prng:
raise unittest.SkipTest("test requires jax_enable_custom_prng")
key = jax.prng.seed_with_impl(jax.prng.rbg_prng_impl, 87)
key = prng.seed_with_impl(prng.rbg_prng_impl, 87)
# Make sure this doesn't crash
pjit(lambda x: x, in_shardings=None, out_shardings=None)(key)