mirror of
https://github.com/ROCm/jax.git
synced 2025-04-14 10:56:06 +00:00
78543f7bb8
Some users of JAX want to use the MLIR dialects defined in jaxlib. In particular, these need to be used by custom lowering rules. Add a semi-public (jax.extend) API to access these, rather than having them use jax._src.lib.mlir. PiperOrigin-RevId: 588448489
16 lines
266 B
ReStructuredText
16 lines
266 B
ReStructuredText
``jax.extend.random`` module
|
|
============================
|
|
|
|
.. 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
|