Remove non-lazy RNG compat mode and flag from flax.

PiperOrigin-RevId: 691326667

CHANGELOG.md

@@ -12,7 +12,7 @@ vNext
 -
 -
 -
--
+- removed FLAX_LAZY_RNG flag support for old non-lazy PRNG derivation mode
 -
 -
 -

flax/configurations.py

@@ -162,9 +162,6 @@ def temp_flip_flag(var_name: str, var_value: bool):
 
 # Flax Global Configuration Variables:
 
-# Whether to use the lazy rng implementation.
-flax_lazy_rng = static_bool_env('FLAX_LAZY_RNG', True)
-
 flax_filter_frames = bool_flag(
   name='flax_filter_frames',
   default=True,

flax/core/scope.py

@@ -38,7 +38,6 @@
 from jax import random, tree_util
 
 from flax import config as config
-from flax import configurations as legacy_config  # only for flax_lazy_rng
 from flax import errors, struct, traceback_util
 from flax.ids import uuid
 from flax.typing import (
@@ -98,11 +97,6 @@ def as_jax_rng(self) -> PRNGKey:
   def create(
     rng: Union['LazyRng', PRNGKey], *suffix: PRNGFoldable
   ) -> 'LazyRng':
-    if not legacy_config.flax_lazy_rng:
-      if isinstance(rng, LazyRng):
-        assert not rng.suffix
-        rng = rng.rng
-      return LazyRng(_legacy_rng_fold_in(rng, suffix), ())
     if isinstance(rng, LazyRng):
       return LazyRng(rng.rng, rng.suffix + suffix)
     else:
@@ -113,22 +107,6 @@ def fold(self):
     return LazyRng(key, ())
 
 
-def _legacy_rng_fold_in(rng: PRNGKey, data: Iterable[PRNGFoldable]) -> PRNGKey:
-  """Legacy RNG folding."""
-  for x in data:
-    if isinstance(x, str):
-      m = hashlib.sha1()
-      m.update(x.encode('utf-8'))
-      d = m.digest()
-      hash_int = int.from_bytes(d[:4], byteorder='big')
-      rng = random.fold_in(rng, jnp.uint32(hash_int))  # type: ignore
-    elif isinstance(x, int):
-      rng = random.fold_in(rng, x)
-    else:
-      raise ValueError(f'Expected int or string, got: {x}')
-  return rng
-
-
 def _fold_in_static(
   rng: PRNGKey, data: typing.Collection[PRNGFoldable]
 ) -> PRNGKey:

tests/run_all_tests.sh

@@ -108,7 +108,6 @@ assert_error="flax is not running on editable mode."
 # env vars must be set after doctest
 export JAX_NUMPY_RANK_PROMOTION=raise
 export FLAX_PROFILE=1
-export FLAX_LAZY_RNG=1
 
 if $RUN_PYTEST; then
   echo "=== RUNNING PYTESTS ==="