Merge pull request #6899 from google:custom-rng

PiperOrigin-RevId: 391913770

jax/__init__.py

@@ -38,7 +38,7 @@
 del _config_module
 
 from ._src.config import (
-  config, enable_checks, check_tracer_leaks, checking_leaks,
+  config, enable_checks, check_tracer_leaks, checking_leaks, enable_custom_prng,
   debug_nans, debug_infs, log_compiles, default_matmul_precision,
   numpy_rank_promotion
 )

jax/_src/config.py

@@ -503,6 +503,15 @@ def update_thread_local_jit_state(**kw):
           'computations. Logging is performed with `absl.logging` at WARNING '
           'level.'))
 
+enable_custom_prng = config.define_bool_state(
+    name='jax_enable_custom_prng',
+    default=False,
+    help=('Enables an internal upgrade that allows one to define custom '
+          'pseudo-random number generator implementations. This will '
+          'be enabled by default in future versions of JAX, at which point '
+          'disabling it will be considered deprecated. In a version '
+          'after that the flag will be removed altogether.'))
+
 hlo_source_file_canonicalization_regex = config.define_string_state(
     name='jax_hlo_source_file_canonicalization_regex',
     default=None,