jax.enable_custom_prng

jax.enable_custom_prng = <jax._src.config._StateContextManager object>

Context manager for jax_enable_custom_prng config option (transient).

Enables an internal upgrade that allows one to define custom pseudo-random number generator implementations. This will be enabled by default in future versions of JAX, at which point all uses of the flag will be considered deprecated (following the API compatibility policy).