jax.experimental.key_reuse module

Experimental Key Reuse Checking

This module contains experimental functionality for detecting re-use of random keys within JAX programs. It is under active development and the APIs here are likely to change. The usage below requires JAX version 0.4.26 or newer.

Key reuse checking can be enabled using the jax_enable_key_reuse_checks configuration:

>>> import jax
>>> jax.config.update('jax_enable_key_reuse_checks', True)
>>> key = jax.random.key(0)
>>> jax.random.normal(key)
Array(-0.20584226, dtype=float32)
>>> jax.random.normal(key)  
Traceback (most recent call last):
 ...
KeyReuseError: Previously-consumed key passed to jit-compiled function at index 0

This flag can also be controlled locally using the jax.enable_key_reuse_checks() context manager:

>>> with jax.enable_key_reuse_checks(False):
...  print(jax.random.normal(key))
-0.20584226

API

reuse_key(key)	Explicitly mark a key as unconsumed.
KeyReuseError