add fold in method

mlcommons · Feb 6, 2024 · f888a99 · f888a99
1 parent 9ebdca7
commit f888a99
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/algorithmic_efficiency/random_utils.py b/algorithmic_efficiency/random_utils.py
@@ -21,6 +21,10 @@
 MAX_INT32 = 2**31
 MIN_INT32 = -MAX_INT32
 
+# SALT constants 
+_SALT1 = np.random.RandomState(seed=5).randint(MIN_INT32, MAX_INT32, dtype=np.int32)
+_SALT2 = np.random.RandomState(seed=6).randint(MIN_INT32, MAX_INT32, dtype=np.int32)
+
 SeedType = Union[int, list, np.ndarray]
 
 
@@ -33,13 +37,11 @@ def _signed_to_unsigned(seed: SeedType) -> SeedType:
     return np.array([s + 2**32 if s < 0 else s for s in seed.tolist()])
 
 
-def _fold_in(seed: SeedType, data: int) -> SeedType:
-  rng_1 = np.random.RandomState(seed=_signed_to_unsigned(seed))
-  new_seed_1 = rng_1.randint(MIN_INT32, MAX_INT32, dtype=np.int32)
-  # Truncate data to 32-bits, since numpy does not support 64-bit ints.
-  rng_2 = np.random.RandomState(seed=_signed_to_unsigned(data) & 0xffffffff)
-  new_seed_2 = rng_2.randint(MIN_INT32, MAX_INT32, dtype=np.int32)
-  return new_seed_1 + new_seed_2
+def _fold_in(seed, data, verbose = True):
+  a = np.random.RandomState(seed=_signed_to_unsigned(seed ^ _SALT1)).randint(MIN_INT32, MAX_INT32, dtype=np.int32)
+  b = np.random.RandomState(seed=_signed_to_unsigned(data ^ _SALT2)).randint(MIN_INT32, MAX_INT32, dtype=np.int32)
+  c = np.random.RandomState(seed=_signed_to_unsigned(a ^ b)).randint(MIN_INT32, MAX_INT32, dtype=np.int32)
+  return c
 
 
 def _split(seed: SeedType, num: int = 2) -> SeedType: