move to count-space-exhaustive partitionable Threefry RNG random_bits

[froystig]

co-authored with @mattjj

Count space exhaustion is achieved by:

• extending iota to arbitrary shapes (not lengths),
• rendering the iota in uint64, then projecting it to a pair of uint32, and
• using the projected pair as the counts for threefry.

Unlike before, this partitionable Threefry random_bits implementation is not bit-equivalent to the original (i.e. without the jax_threefry_partitionable config enabled). The two produce different random bits from the same RNG key.