jax-keyguard

The problem

JAX has no global random state. Every random function takes an explicit key, and the same key always gives the same numbers. To get new numbers you have to split the key and use each part once. Forgetting to split is easy, and nothing warns you: the code runs and the output looks plausible.

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

def init(key):
    w_key, b_key = jax.random.split(key)
    w = jax.random.normal(w_key, (3, 3))
    b = jax.random.normal(b_key, (3,))
    return w, b

Why JAX's own check isn't enough

JAX 0.4.26 added a jax_debug_key_reuse option. I tested what it actually catches:

PRNGKey is the older kind of key, but it still works in the latest JAX and is used in most existing code and tutorials. jax-keyguard checks both kinds on any JAX version from 0.4.18 on, and its static checker doesn't depend on the JAX version at all.

Two ways to check

1. Static checker: reads the code, never runs it

Run it on a file or a directory, or use it as a flake8 plugin. Because the code isn't executed, it also checks functions under jax.jit, and it can run in CI with no GPU and no test data. Given a file containing the buggy init above and a jitted dropout function that splits its key and then uses the original anyway:

$ jax-keyguard src/
src/model.py:6:27: KG001 PRNG key 'key' reused: already consumed by jax.random.normal on line 5; split it first
src/model.py:13:33: KG001 PRNG key 'key' reused: already consumed by jax.random.split on line 12; split it first

It exits with status 1 when it finds something, so a CI job fails on reuse. Lines can be silenced with # noqa.

2. Runtime check: raises an error the moment a key is reused

Add --keyguard to pytest, or call jax_keyguard.enable() in a script. The error names the line that used the key first:

$ pytest --keyguard
E   jax_keyguard.runtime.KeyReuseError: jax.random.normal received a PRNG key that was already consumed by jax.random.normal at tests/test_model.py:5. Split the key and use each subkey only once.
FAILED tests/test_model.py::test_init

This catches cases the static checker can't see, such as keys passed through your own helper functions or stored in containers. The two complement each other.

Found in JAX's own repository

I ran the static checker on the JAX repository (1,075 Python files, about 28 seconds). In one of JAX's example programs, jax2tf/examples/mnist_lib.py, it reports this:

rng = jax.random.PRNGKey(0)
params = [(param_scale * jax.random.normal(rng, (m, n)),
           param_scale * jax.random.normal(rng, (n,)))
          for m, n, in zip(layer_sizes[:-1], layer_sizes[1:])]

mnist_lib.py:173:48: KG001 PRNG key 'rng' reused: already consumed by jax.random.normal in a previous iteration; split it first
mnist_lib.py:174:48: KG001 PRNG key 'rng' reused: already consumed by jax.random.normal on line 173; split it first

The weights and the biases of every layer are drawn from the same key, so every layer starts with the same values. The code uses PRNGKey, so JAX's own jax_debug_key_reuse would not have caught it even if it had been turned on. Most of the other reports were in JAX's test files, some of them on purpose, such as the tests for JAX's own reuse check.

How the static checker works

It parses each file with Python's ast module and follows every variable that holds a key through each function, the way a compiler tracks values:

• What counts as a key: anything returned by key, PRNGKey, split or fold_in, copies like k2 = key, loop variables over split keys, and function parameters that are passed to jax.random.
• Branches: the two sides of an if are checked separately and then merged, so using the key once in each branch is fine.
• Loops: the loop body is re-analyzed until nothing changes, which finds a key used again on the next iteration. break, continue, try/finally, with and match are handled too.
• Imports: import jax.random as jr, from jax import random and from jax.random import normal are all recognized.

What it deliberately leaves out

A checker that raises many false alarms gets turned off, so it only reports what it can be sure of. It doesn't follow keys into your own helper functions, because those often call fold_in, which is safe. It also doesn't track keys stored in attributes or lists. The runtime check covers these cases for code that runs.

Testing and results

Test suite

101 tests, covering both checkers and the pytest plugin

Python versions

3.10, 3.11, 3.12 and 3.13

JAX versions

0.4.18 up to 0.12

Runtime cost

About twice as slow for random calls with checking on, so it's meant for tests and debugging rather than full training runs

Dependencies

Only JAX. The static checker doesn't even import it.

Try it

pip install jax-keyguard

jax-keyguard src/          # static check
pytest --keyguard          # runtime check in tests