liveBook · Manning

Fold_in()

As explained in Deep Learning with JAX

A JAX random utility that combines an existing key with integer data to derive a new pseudo-random key.

Need a distinct JAX random key for each loop iteration without allocating all keys in advance? random.fold_in() combines an existing key with integer data and returns a statistically safe derived key. JAX keys are explicit state: reuse the same key when identical output is intended, and discard keys after consuming them. section 9.2.3 section 9.4

Derive keys inside a loop

key = random.PRNGKey(42) for i in range(5): new_key = random.fold_in(key, i) vals = random.normal(new_key, shape=(3, 5)) # use vals Each iteration produces a different key and random array. The integer may also represent stable data such as a layer identifier; convert strings to integers first, for example with a deterministic SHA-1-based hash. section 9.2.3

Listing 9.18 Using fold_in() to generate new keys
key = random.PRNGKey(42)                                      #1
 
for i in range(5):
  new_key = random.fold_in(key, i)                            #2
  print(new_key)
  vals = random.normal(new_key, shape=(3,5))                  #3
  # do something with values
 
>>> [1832780943  270669613]
>>> [  64467757 2916123636]
>>> [2465931498  255383827]
>>> [3134548294  894150801]
>>> [2954079971 3276725750]

Shows fold_in() deriving a fresh key from a loop index and using it with random.normal(). section 9.2.3

How Deep Learning with JAX covers it

How deep it goes
Explains key derivation, integer-only input, string hashing, and key-reuse rules, but does not detail the underlying PRNG algorithm.
Examples
Demonstrates loop indices and hashed neural-network layer names as data folded into a JAX key.
What it assumes
Assumes basic JAX random keys and random number generation while teaching explicit key consumption and integer conversion.
What you can do after
Apply it via JAX random.fold_in() to derive per-iteration or per-label keys for random number generation.
How it connects
Connects fold_in() with JAX PRNG keys, split(), random.normal(), explicit state, and NumPy’s contrasting random-state behavior.
Cover of Deep Learning with JAX
liveBookmanning.com
Read Deep Learning with JAX on liveBook Get the book at manning.com

This explainer was generated by Manning's knowledge graph from the text of Deep Learning with JAX. The section links point to the parts of the book it draws on.