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

Merge pull request #10362 from JeppeKlitgaard:fix-random-reexport

Browse Source 
PiperOrigin-RevId: 442936248
This commit is contained in:
jax authors 2022-04-19 16:04:31 -07:00
commit 88ac4edf63
1 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
jax/random.py
View File

@ -102,18 +102,18 @@ The possible reasons not use the default RNG are:
Here is a short summary:
.. table::
.. table::
:widths: auto
================================= ================= === ==========
Property ThreeFry, default rbg unsafe_rbg
================================= ================= === ==========
Fast on TPU ✅ ✅
always correct w/ scan ✅ ✅
always correct w/ remat ✅ ✅
identical across CPU/GPU/TPU ✅ ✅
identical across JAX/XLA versions ✅
identical across shardings ✅
always correct w/ scan ✅ ✅
always correct w/ remat ✅ ✅
identical across CPU/GPU/TPU ✅ ✅
identical across JAX/XLA versions ✅
identical across shardings ✅
================================= ================= === ==========
"""
@ -122,7 +122,8 @@ Here is a short summary:
# TODO(frostig): replace with KeyArray from jax._src.random once we
# always enable_custom_prng
from jax._src.prng import PRNGKeyArray as KeyArray
from jax._src.prng import PRNGKeyArray
KeyArray = PRNGKeyArray
from jax._src.random import (
PRNGKey as PRNGKey,