Replace iterative RNG split and carry with jax.random.fold_in

[junpenglao]

Following new recommendation on RNG usage within a loop.

import jax
key = jax.random.key(42)

Main pattern: Instead of doing

for step in range(10000000):
  key, subkey = jax.random.split(key, 2)
  values = jax.random.normal(subkey, [3, 5])
  # ... use values ...

do

for step in range(10000000):
  subkey = jax.random.fold_in(key, step)
  values = jax.random.normal(subkey, [3, 5])
  # ... use values ...

The main concern here is that there is an unnecessary statistical inefficiency using the split-then-carry within a loop. By doing splitting, random number is generated by iterating the hash on keys at a fixed message, so the resulting sequence might have a shorter period than you expect.

Also change jax.tree_map to jax.tree.map

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 98.87%. Comparing base (7cf4f9d) to head (b1ca8a1).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #656      +/-   ##
==========================================
- Coverage   98.87%   98.87%   -0.01%     
==========================================
  Files          59       59              
  Lines        2745     2744       -1     
==========================================
- Hits         2714     2713       -1     
  Misses         31       31              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.