jax
0fa13e32 - Add philox4x32 PRNG impl

Commit
59 days ago
Add philox4x32 PRNG impl This can be used directly with `jax.random.key`: ```python >>> import jax >>> key = jax.random.key(0, impl='philox4x32') >>> jax.random.uniform(key) Array(0.8307482, dtype=float32) ``` Philox4x32 has a 64-bit key space and a 128-bit counter space (compared to threefry2x32 64-bit key space and 64-bit counter space, and philox2x32 32-bit key space and 64-bit counter space). Like philox2x32, it uses operations that we may be able to lower more efficiently on accelerators (optimizations to be done later). The increased state size is useful for context where many samples will be generated. PiperOrigin-RevId: 917843748
Author
Jake VanderPlas
Parents
Loading