Details on the threefry2x32 implemenetation

[botev-openai]

I could not find any details on the threefry2x32 implementation, in particular Jax threefry2x32 primitive takes two uint32 keys and two counters (outputs of iota2x32) and combines them to produce 2 outputs of the same shape as the counters, which are then used together to produce the two numbers for each Key.
My question is that the usual Threefry2x32 takes one key and one counter, not two. Are there any docs on how are these used in the actual threefry implementation?

[jakevdp]

JAX's PRNG desing is discussed at JAX PRNG Design, and references the paper Parallel numbers: easy as 1, 2, 3 (pdf). JAX implements the 2x32 variant, in which the key and the produced samples are represented by pairs of 32-bit integers.