Roy Frostig 16d082b002 [jex] replace extend.random.PRNGImpl with extend.random.define_prng_impl ...

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

2023-10-31 17:21:54 -07:00

..

__init__.py

JEX: add jax.extend.source_info_util

2023-10-13 09:36:00 -07:00

core.py

JEX: add jex.abstract_arrays.array_types

2023-09-19 11:37:05 -07:00

linear_util.py

JEX: move jax.linear_util to jax.extend.linear_util

2023-08-30 18:32:12 -07:00

random.py

[jex] replace extend.random.PRNGImpl with extend.random.define_prng_impl

2023-10-31 17:21:54 -07:00

source_info_util.py

export jax.extend.source_info_util.current

2023-10-13 12:31:11 -07:00