Two words and a hash

Ask a key to print itself and it shows you both halves of what it is. The repr names an array of shape () with dtype key<fry>, and underneath it prints two unsigned 32-bit words: 0 and 0. A seed of zero becomes the pair (0, 0) literally, and jax.random.key_data hands those words back as a plain uint32 array any time you want to look.

The older constructor returns the same two words with the type peeled off. jax.random.PRNGKey(0) is uint32 of shape (2,), and it equals key_data of the typed key entry for entry. Nothing about the randomness differs between the two calls. What differs is whether the array carries its own type.

That type is a field on the value, not a convention you hold in your head. jax.random.key_impl reports threefry2x32 for this key, so the key carries its own generator. A uint32 key carries no such tag, and the module docs are direct about the consequence: legacy keys do not carry information about the RNG implementation, so a global configuration setting decides which algorithm runs.

import jax

k = jax.random.key(0)
print(repr(k))
print(k.dtype, k.shape, jax.random.key_impl(k))
print(jax.random.key_data(k))
print(jax.random.PRNGKey(0), jax.random.PRNGKey(0).dtype, jax.random.PRNGKey(0).shape)

# Array((), dtype=key<fry>) overlaying:
# [0 0]
# key<fry> () threefry2x32
# [0 0]
# [0 0] uint32 (2,)

Add one to a typed key and JAX refuses, with add does not accept dtypes key<fry>, int32. There is no arithmetic on a key, because arithmetic on a key is never what anyone meant. The legal moves are all key moves: split it, fold an integer into it, hand it to a sampler, or read its raw words out.

Do the same thing to a legacy key and it goes through. jax.random.PRNGKey(0) + 1 returns [1 1], an ordinary uint32 array that every sampling function will accept as a key. The result is a seed nobody chose, produced by a line that looks like a typo and runs like an instruction.

Getting back to plain integers is still one call, for the times a serializer needs them. key_data pulls the two words out, wrap_key_data puts them back, and the round trip returns the pair unchanged.

import jax

k = jax.random.key(0)
try:
    k + 1
except TypeError as err:
    print(err)
print(jax.random.PRNGKey(0) + 1)
print(jax.random.key_data(jax.random.wrap_key_data(jax.random.key_data(k))))

# add does not accept dtypes key<fry>, int32.
# [1 1]
# [0 0]

Trace a call to jax.random.bits and the entire program is one equation. The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → holds a single primitive, random_bits, and the requested shape sits inside the brackets as an attribute rather than flowing in as a value. The shape is not packaging around the draw. It is an input to it.

Under that primitive, threefry2x32 hashes the key against a count array, and for a 32-bit request that count array is an iota: 0, 1, 2, and so on, one entry per word. Word i of the output depends on the key and on i, and on nothing else in the array. That independence is why a draw parallelizes and why no element waits on its neighbour.

Two requests with the same element count therefore line up exactly. A (2, 2) normal draw reshaped to (4,) is elementwise equal to the (4,) draw, because the reshape happens after the same iota was hashed. Requests with different element counts are a different story, and the next section is the mechanism behind it.

import jax

k = jax.random.key(0)
print(jax.make_jaxpr(lambda key: jax.random.bits(key, (4,)))(k))
print(jax.random.bits(k, (4,)))
print(jax.random.normal(k, (2, 2)).reshape(4))
print(jax.random.normal(k, (4,)))

# { lambda ; a:key<fry>[]. let
#     b:u32[4] = random_bits[bit_width=32 shape=(4,)] a
#   in (b,) }
# [4146024105  967050713 2718843009 1272950319]
# [ 1.8160863  -0.75488514  0.33988908 -0.53483534]
# [ 1.8160863  -0.75488514  0.33988908 -0.53483534]

Threefry hashes two words at a time, so the count array is split down the middle and the two halves go in as a pair. An odd-sized request gets a zero appended first and the extra output dropped at the end. Four lines of threefry_2x32 do all of it.

Run the halving by hand for a four-word request and a three-word one. The four-word counts split into [0, 1] and [2, 3], so position 0 is hashed against 2 and position 1 against 3. The three-word counts pad to [0, 1] and [2, 0], so position 0 is still hashed against 2 while position 1 now meets a zero. That is exactly the pattern the run shows: entries 0 and 2 match across the two requests, entry 1 does not.

Which means the overlap you can see is an artifact of one implementation choice, not a rule. Under the partitionable flag that jax 0.5.0 turns on by default, the counter is built per element instead of per half, and a three-word request becomes a strict prefix of a four-word one. Do not build anything on either behaviour; build on the shape being part of the draw.

The same key with a different shape is a different draw.

  if odd_size:
    x = list(jnp.split(jnp.concatenate([count.ravel(), np.uint32([0])]), 2))
  else:
    x = list(jnp.split(count.ravel(), 2))

  x = threefry2x32_p.bind(key1, key2, x[0], x[1])
  out = jnp.concatenate(x)
  assert out.dtype == np.uint32
  return lax.reshape(out[:-1] if odd_size else out, count.shape)

# >>> jax.random.bits(k, (4,))
# Array([4146024105,  967050713, 2718843009, 1272950319], dtype=uint32)
# >>> jax.random.bits(k, (3,))
# Array([4146024105, 1351547692, 2718843009], dtype=uint32)
# >>> jax.random.normal(k, (4,))
# Array([ 1.8160863 , -0.75488514,  0.33988908, -0.53483534], dtype=float32)
# >>> jax.random.normal(k, (3,))
# Array([ 1.8160863 , -0.48262316,  0.33988908], dtype=float32)

Three lines implement splitting, and none of them is new machinery. _threefry_split_original takes an iota twice as long as the number of children you asked for, hashes it with the same threefry call every sampler uses, and reshapes the result into pairs. A child key is two random words wearing the key type.

The session proves it end to end. The eight words behind split(k, 4) are elementwise equal to bits(k, (8,)), taken from the same key. Splitting is not a separate source of randomness sitting next to sampling. It is sampling, with the output labelled as keys.

The count you pass therefore changes every child, not just the extra ones. split(k, 2) gives [4146024105, 967050713] and [2718843009, 1272950319], which are the first four words of the eight above only because 2 happens to divide 4 evenly here; ask for a different width and the halving from the previous section re-pairs everything. Treat the number of children as part of the draw, the same way you treat a sampler's shape.

@partial(jit, static_argnums=(1,), inline=True)
def _threefry_split_original(key, shape) -> typing.Array:
  num = math.prod(shape)
  counts = lax.iota(np.uint32, num * 2)
  return lax.reshape(threefry_2x32(key, counts), (*shape, 2))

# >>> jax.random.key_data(jax.random.split(k, 4)).reshape(-1)
# Array([2285895361, 1501764800, 1518642379, 4090693311,  433833334,
#        4221794875,  839183663, 3740430601], dtype=uint32)
# >>> jax.random.bits(k, (8,))
# Array([2285895361, 1501764800, 1518642379, 4090693311,  433833334,
#        4221794875,  839183663, 3740430601], dtype=uint32)
# >>> jax.random.key_data(jax.random.split(k, 2))
# Array([[4146024105,  967050713],
#        [2718843009, 1272950319]], dtype=uint32)

Folding an integer in is the same hash against a different counter. _threefry_fold_in calls threefry_2x32 with threefry_seed(data), and threefry_seed of a small integer is the pair (0, i), which is exactly what PRNGKey(i) prints. So fold_in(k, 3) hashes the key against the counts (0, 3), while split(k, n) hashes it against an iota. Different counters, one hash, no shared state between them.

The counter also explains a property worth banking. fold_in's counts depend on i alone, never on how many siblings exist or on which of them ran first. Ask for the key of item 3 and you get the same two words whether you asked for items 0 through 7 or for item 3 by itself.

One measurement makes the difference concrete. Turning jax_threefry_partitionable on, the flag jax 0.5.0 ships enabled, changes the words that split returns; the words that fold_in returns are identical to the run below, for every index 0 through 7 on this machine. Draws taken from those keys still move, because random_bits changed too. The key derivation is what stayed put.

def threefry_fold_in(key: typing.Array, data: typing.Array) -> typing.Array:
  assert not data.shape
  return _threefry_fold_in(key, jnp.uint32(data))

@jit
def _threefry_fold_in(key, data):
  return threefry_2x32(key, threefry_seed(data))

# >>> jax.random.PRNGKey(3)
# Array([0, 3], dtype=uint32)
# >>> jax.random.key_data(jax.random.fold_in(k, 3))
# Array([2467461003, 3840466878], dtype=uint32)