doctest: avoid modifying global flag state

Author: jakevdp

jax/experimental/key_reuse/__init__.py

@@ -16,26 +16,27 @@
 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.
+This module contains **experimental** functionality for detecting reuse 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::
+Key reuse checking can be enabled using the ``jax_enable_key_reuse_checks`` configuration.
+This can be set globally using::
+
+  >>> jax.config.update('jax_enable_key_reuse_checks', True)  # doctest: +SKIP
+
+Or it can be enabled locally with the :func:`jax.enable_key_reuse_checks` context manager.
+When enabled, using the same key twice will result in a :class:`~jax.errors.KeyReuseError`::
 
   >>> 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)  # doctest: +IGNORE_EXCEPTION_DETAIL
+  >>> with jax.enable_key_reuse_checks(True):
+  ...   key = jax.random.key(0)
+  ...   val1 = jax.random.normal(key)
+  ...   val2 = jax.random.normal(key)  # doctest: +IGNORE_EXCEPTION_DETAIL
   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 :func:`jax.enable_key_reuse_checks`
-context manager::
-
-  >>> with jax.enable_key_reuse_checks(False):
-  ...  print(jax.random.normal(key))
-  -0.20584226
+The key reuse checker is currently experimental, but in the future we will likely
+enable it by default.
 """