Pseudorandom number generation in JAX

Recently I’ve been working on a statistics library in Typst called distro*More to come soon. Typst is architected with an incremental compiler, and pairs it with pure functions*Almost all functions in Typst are pure, with a few exceptions to facilitate snappy real-time compilation. The problem is, what to do about random number generation when your functions are supposed to be referentially transparent? Well, thankfully that’s a solved problem with a neat solution. For instance, JAX relies on an explicit ‘key’ value, that is consumed by the PRNG, and returns an updated state that is then passed back into the PRNG to generate the next random sample. PRNG-generated sequences are not truly random because they are actually determined by their initial value, which is typically referred to as the seed, and each step of random sampling is a deterministic function of some state that is carried over from a sample to the next. Importantly, they warn you to never re-use the same key, else: The…