Reuse is silent

Initialize two weight matrices from one key and they come back identical. Not similar, not correlated: the same nine numbers, twice, because a key plus a shape determines the bits and nothing else about the two calls differed.

No check fires. The shapes are right, the dtypes are right, the values are plausible draws from a normal, and a model built on them will train. It started from a smaller family of initializations than the code claims, and no output says so.

Which is why this one has no exhibit next door. The museum's cases, inplace-assign and tracer-bool among them, all end in a message you can paste into a search box. There is no message to quote here, so the rest of this lesson is about the symptoms you can measure instead.

import jax
import jax.numpy as jnp

key = jax.random.key(0)
w1 = jax.random.normal(key, (3, 3))
w2 = jax.random.normal(key, (3, 3))
print(bool(jnp.all(w1 == w2)))
print(w1)

# True
# [[-0.3721109   0.26423115 -0.18252768]
#  [-0.7368197   0.44973662 -0.1521442 ]
#  [-0.67135346 -0.5908641   0.73168886]]

Correlated is the usual word for what reuse does, and it undersells the damage. Draw ten thousand normals from a key, draw ten thousand more from the same key, and the correlation comes back 1.0 to fourteen decimal places while the largest elementwise difference is exactly 0.0. The second draw is not a correlated sample of the first. It is the first.

Split the key and the same measurement lands where independence lands. The two children give -0.0055 over ten thousand samples, and a parent's draw measured against one of its children gives -0.0033. Both are the size of sampling noise at n = 10000, and neither is evidence of anything except that the keys were different.

So the test for reuse is not statistical. It is equality, and it is cheap: assert that no two draws you meant to be independent are elementwise equal.

import jax
import jax.numpy as jnp
import numpy as np

key = jax.random.key(0)
k1, k2 = jax.random.split(key)
x = jax.random.normal(key, (10000,))
y = jax.random.normal(key, (10000,))
u = jax.random.normal(k1, (10000,))
v = jax.random.normal(k2, (10000,))
corr = lambda a, b: float(np.corrcoef(np.asarray(a), np.asarray(b))[0, 1])
print(corr(x, y), float(jnp.max(jnp.abs(x - y))))
print(corr(u, v))
print(corr(x, u))

# 0.9999999999999998 0.0
# -0.0054627058855592735
# -0.0033143140158316963

Averaging is where reuse turns into a number somebody eventually notices. Eight independent draws of 2048 normals, averaged elementwise, give a mean array with standard deviation 0.352, which is 1 over the square root of 8, or 0.3536, as it should be. Eight draws from one key, averaged the same way, give 0.978.

Nothing got slower and nothing warned. The estimator simply stopped converging at the rate the code was written to assume, which is the shape this bug takes in a Monte Carlo estimate, a dropout ensemble, an augmentation pipeline, or a bootstrap.

If you only remember one diagnostic from this lesson, remember this one. Averaging N samples should shrink the spread by a factor near the square root of N. When it does not, look at the keys before you look at the model.

import jax
import jax.numpy as jnp

key = jax.random.key(0)
keys = jax.random.split(key, 8)
independent = jax.vmap(lambda k: jax.random.normal(k, (2048,)))(keys).mean(0)
reused = jnp.stack([jax.random.normal(key, (2048,)) for _ in range(8)]).mean(0)
print(float(independent.std()), float(reused.std()), float(1 / jnp.sqrt(8.0)))

# 0.3524690568447113 0.9781718850135803 0.3535533845424652

The idiom in the jax.random docstring is one line long and every character of it carries weight: key, subkey = jax.random.split(key). The left-hand key is a rebinding, and dropping it turns the loop into three copies of one step.

The stale loop below draws the integers 8, 45, 78, 7 three times over. It calls split on every iteration, so the code looks like it is doing the right thing; what it never does is move the name forward, so the same parent produces the same child each time round.

A key is a value like any other, which means Python will not stop you from reading a stale one. The discipline is a naming discipline: once a key has been split, the parent name should point at the new key, and a key you have already sampled from should never appear again.

>>> seed = 1701
>>> num_steps = 100
>>> key = jax.random.key(seed)
>>> for i in range(num_steps):
...   key, subkey = jax.random.split(key)
...   params = compiled_update(subkey, params, next(batches))

import jax
import jax.numpy as jnp

key = jax.random.key(0)

k = key
stale = []
for _ in range(3):
    stale.append(jax.random.randint(jax.random.split(k)[1], (4,), 0, 100))

k = key
fresh = []
for _ in range(3):
    k, sub = jax.random.split(k)
    fresh.append(jax.random.randint(sub, (4,), 0, 100))

print(jnp.stack(stale))
print(jnp.stack(fresh))

# [[ 8 45 78  7]
#  [ 8 45 78  7]
#  [ 8 45 78  7]]
# [[ 8 45 78  7]
#  [81 73 74 37]
#  [43 35 56 43]]

JAX ships a checker for exactly this failure, switched off by default. Wrap the code in jax.debug_key_reuse(True) and a consumed key becomes a KeyReuseError, with two messages depending on where the reuse happened: In pjit, argument 0 is already consumed. when both draws are inside one jitted function, and Previously-consumed key passed to jit-compiled function at index 0 when a key crosses the boundary twice.

Repeating a draw on purpose is still allowed, as long as you say so. jax.random.clone returns a key the checker treats as fresh, and the values that come out match the original exactly, which is usually the reason you wanted it: replaying the same noise for a comparison.

The checker has one blind spot, and it is the sharpest argument in the arc for typed keys. It tracks key<fry> values, so a program built on jax.random.PRNGKey sails through: two draws from the same uint32 key under the checker return identical arrays and raise nothing. The module's own documentation calls the checker experimental and says that in the future we will likely enable it by default, which is a good reason to run it now and find out what it says about your code.

import jax
import jax.numpy as jnp

def init(key):
    return jax.random.normal(key, (3, 3)), jax.random.normal(key, (3,))

with jax.debug_key_reuse(True):
    try:
        jax.jit(init)(jax.random.key(0))
    except jax.errors.KeyReuseError as err:
        print(err)
    key = jax.random.key(0)
    a = jax.random.normal(key, (2,))
    try:
        jax.random.normal(key, (2,))
    except jax.errors.KeyReuseError as err:
        print(str(err).splitlines()[0])
    print(bool(jnp.all(a == jax.random.normal(jax.random.clone(key), (2,)))))

    raw = jax.random.PRNGKey(0)
    print(bool(jnp.all(jax.random.normal(raw, (2,)) == jax.random.normal(raw, (2,)))))

# In pjit, argument 0 is already consumed.
# See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.KeyReuseError
# Previously-consumed key passed to jit-compiled function at index 0
# True
# True