flax
0d69f1a2 - Remove unnecessary splitting of rngs in lifting logic.

Commit
4 years ago
Remove unnecessary splitting of rngs in lifting logic. Previosuly we split every rng that was lifted. Now the rng_counters are shared between the lifted and unlifted scope. This is valid because make_rng can only be called a static number of times and lifted loops/vmap split the rngs manually already
Author
Parents
Loading