Allow pluggable PRNG implementations

[hawkinsp]

Currently we use ThreeFry2x32 as our PRNG, but it has an uncomfortably small key space (2^64).

An easy way to improve this would be to use the 4x32 variant of ThreeFry instead which has a key space of size 2^128.

[jakevdp]

This sounds fun...

For backward compatibility, do you think it might make sense to implement this beside the current approach (e.g. with jax.random.PRNGKey4() or something similar?)

[jakevdp]

Or maybe better, a keysize keyword to PRNGKey that defaults to 2, and can optionally be set to 4, with the intention of eventually changing the default?

[hawkinsp]

I was thinking that perhaps we would do something like PRNGKey(..., algorithm=THREEFRY4x32) to allow for the possiblity we might want to support other PRNGs.

[NeilGirdhar]

#2294 might be related? It would allow you to make this change more transparently since no code would assume a key size; it would be hidden behind the opaque object.

[hawkinsp]

Generalizing this issue: we should allow pluggable PRNG bit generators (e.g., Philox on GPU).

[froystig]

#6899 allows for custom PRNG implementations. Filed #7676 to follow up and use it to add a wide PRNG.