jax
16d082b0 - [jex] replace `extend.random.PRNGImpl` with `extend.random.define_prng_impl`

Commit
2 years ago
[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
Author
Committer
Parents
Loading