Merge pull request #20370 from jakevdp:key-reuse-flag

PiperOrigin-RevId: 618179711

Author: jax authors

jax/__init__.py

@@ -47,7 +47,7 @@
 from jax._src.config import (
   config as config,
   enable_checks as enable_checks,
-  enable_key_reuse_checks as enable_key_reuse_checks,
+  debug_key_reuse as debug_key_reuse,
   check_tracer_leaks as check_tracer_leaks,
   checking_leaks as checking_leaks,
   enable_custom_prng as enable_custom_prng,

jax/_src/config.py

@@ -213,7 +213,7 @@ def trace_context():
           softmax_custom_jvp.value,
           enable_memories.value,
           disable_jit.value,
-          enable_key_reuse_checks.value,
+          debug_key_reuse.value,
           jax_xla_profile_version.value,
           # Technically this affects jaxpr->stablehlo lowering, not tracing.
           hlo_source_file_canonicalization_regex.value)
@@ -930,8 +930,8 @@ def update_thread_local_jit_state(**kw):
     default=False,
     help='Turn on invariant checking for JAX internals. Makes things slower.')
 
-enable_key_reuse_checks = define_bool_state(
-    name='jax_enable_key_reuse_checks',
+debug_key_reuse = define_bool_state(
+    name='jax_debug_key_reuse',
     default=False,
     help=('Turn on experimental key reuse checking. With this configuration enabled,'
           ' typed PRNG keys (i.e. keys created with jax.random.key()) will have their'

jax/_src/core.py

@@ -2861,7 +2861,7 @@ def ctx_factory():
     raise JaxprTypeError(msg) from None
 
   # Run key reuse checker after validating jaxpr:
-  if config.enable_key_reuse_checks.value:
+  if config.debug_key_reuse.value:
     # Import here to avoid circular imports
     from jax.experimental.key_reuse._core import check_key_reuse_jaxpr  # pytype: disable=import-error
     check_key_reuse_jaxpr(jaxpr)

jax/_src/errors.py

@@ -661,12 +661,12 @@ def __init__(self, msg: str):
 class KeyReuseError(JAXTypeError):
   """
   This error occurs when a PRNG key is reused in an unsafe manner.
-  Key reuse is checked only when `jax_enable_key_reuse_checks` is
+  Key reuse is checked only when `jax_debug_key_reuse` is
   set to `True`.
 
   Here is a simple example of code that would lead to such an error::
 
-    >>> with jax.enable_key_reuse_checks(True):  # doctest: +SKIP
+    >>> with jax.debug_key_reuse(True):  # doctest: +SKIP
     ...   key = jax.random.key(0)
     ...   value = jax.random.uniform(key)
     ...   new_value = jax.random.uniform(key)

jax/_src/pjit.py

@@ -236,7 +236,7 @@ def _get_fastpath_data(
       # no ref state effects
       and not any(isinstance(e, RefEffect) for e in effects)
       # no prng reuse checking
-      and not (config.enable_key_reuse_checks.value and any(
+      and not (config.debug_key_reuse.value and any(
         hasattr(arg, 'dtype') and dtypes.issubdtype(arg.dtype, dtypes.prng_key)
         for arg in (*args_flat, *out_flat)))
       )
@@ -1150,7 +1150,7 @@ def _create_pjit_jaxpr(fun, in_type, debug_info, out_paths, ignored_inline):
   if not config.dynamic_shapes.value and not attrs_tracked:
     jaxpr = jaxpr_debug_info(jaxpr, debug_info, out_paths())
 
-  if config.enable_key_reuse_checks.value:
+  if config.debug_key_reuse.value:
     # Import here to avoid circular imports
     from jax.experimental.key_reuse._core import check_key_reuse_jaxpr
     check_key_reuse_jaxpr(jaxpr)

jax/experimental/jax2tf/tests/jax_primitives_coverage_test.py

@@ -42,7 +42,7 @@
 
 
 @jtu.with_config(jax_legacy_prng_key='allow',
-                 jax_enable_key_reuse_checks=False)
+                 jax_debug_key_reuse=False)
 class JaxPrimitiveTest(jtu.JaxTestCase):
 
   # This test runs for all primitive harnesses. For each primitive "xxx" the

jax/experimental/jax2tf/tests/tf_test_util.py

@@ -158,7 +158,7 @@ def ComputeTfValueAndGrad(tf_f: Callable, tf_args: Sequence,
 @jtu.with_config(jax_numpy_rank_promotion="allow",
                  jax_numpy_dtype_promotion='standard',
                  jax_legacy_prng_key="allow",
-                 jax_enable_key_reuse_checks=False)
+                 jax_debug_key_reuse=False)
 class JaxToTfTestCase(jtu.JaxTestCase):
   # We want most tests to use the maximum available version, from the locally
   # installed tfxla module and export.

jax/experimental/key_reuse/__init__.py

@@ -20,16 +20,16 @@
 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.
+Key reuse checking can be enabled using the ``jax_debug_key_reuse`` configuration.
 This can be set globally using::
 
-  >>> jax.config.update('jax_enable_key_reuse_checks', True)  # doctest: +SKIP
+  >>> jax.config.update('jax_debug_key_reuse', True)  # doctest: +SKIP
 
-Or it can be enabled locally with the :func:`jax.enable_key_reuse_checks` context manager.
+Or it can be enabled locally with the :func:`jax.debug_key_reuse` context manager.
 When enabled, using the same key twice will result in a :class:`~jax.errors.KeyReuseError`::
 
   >>> import jax
-  >>> with jax.enable_key_reuse_checks(True):
+  >>> with jax.debug_key_reuse(True):
   ...   key = jax.random.key(0)
   ...   val1 = jax.random.normal(key)
   ...   val2 = jax.random.normal(key)  # doctest: +IGNORE_EXCEPTION_DETAIL

jax/experimental/key_reuse/_core.py

@@ -530,7 +530,7 @@ def _remat_key_type_signature(eqn):
 def key_reuse_impl_rule(prim, original_rule):
   @wraps(original_rule)
   def key_reuse_impl(*args, **kwargs):
-    if config.enable_key_reuse_checks.value:
+    if config.debug_key_reuse.value:
       if prim == pjit.pjit_p:
         funcname = "jit-compiled function"
         jaxpr = kwargs['jaxpr'].jaxpr

tests/batching_test.py

@@ -962,7 +962,7 @@ def body_fn(uk):
       u, _ = lax.while_loop(lambda uk: uk[0] > 0.5, body_fn, (1., key))
       return u
 
-    with jax.enable_key_reuse_checks(False):
+    with jax.debug_key_reuse(False):
       print(vmap(f)(random.split(random.PRNGKey(0), 2)))  # no crash
 
   def testEmptyTuples(self):