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

Merge pull request #17337 from jakevdp:jax-extend-doc

Browse Source 
PiperOrigin-RevId: 561101464
This commit is contained in:
jax authors 2023-08-29 12:11:23 -07:00
commit ff5b480c6b
2 changed files with 19 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
docs/jax.extend.rst
View File

@ -2,3 +2,20 @@
=====================
.. automodule:: jax.extend
``jax.extend.random``
---------------------
.. automodule:: jax.extend.random
.. autosummary::
:toctree: _autosummary
PRNGImpl
seed_with_impl
threefry2x32_p
threefry_2x32
threefry_prng_impl
rbg_prng_impl
unsafe_rbg_prng_impl
wrap_key_data

2
docs/jax.random.rst
View File

@ -16,6 +16,8 @@ List of Available Functions
:toctree: _autosummary
PRNGKey
key
key_data
ball
bernoulli
beta