mirror of
https://github.com/ROCm/jax.git
synced 2025-04-18 04:46:06 +00:00
16d082b002
Instead of exposing a constructor, only expose a function that returns an opaque object representing the defined implementation. This result can still be passed to `jax.random.key` and `wrap_key_data`. PiperOrigin-RevId: 578349699
38 lines
577 B
ReStructuredText
38 lines
577 B
ReStructuredText
``jax.extend`` module
|
|
=====================
|
|
|
|
.. automodule:: jax.extend
|
|
|
|
``jax.extend.linear_util``
|
|
--------------------------
|
|
|
|
.. automodule:: jax.extend.linear_util
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
StoreException
|
|
WrappedFun
|
|
cache
|
|
merge_linear_aux
|
|
transformation
|
|
transformation_with_aux
|
|
wrap_init
|
|
|
|
|
|
``jax.extend.random``
|
|
---------------------
|
|
|
|
.. automodule:: jax.extend.random
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
define_prng_impl
|
|
seed_with_impl
|
|
threefry2x32_p
|
|
threefry_2x32
|
|
threefry_prng_impl
|
|
rbg_prng_impl
|
|
unsafe_rbg_prng_impl
|