What's the best way to index into a PRNG sequence?

[NeilGirdhar]

Given a sequence as A below, how do I get the ith element of the sequence without resplitting and indexing? One idea I had was to use fold-in, and then I can reproduce the ith element by folding in i. Is there any downside to folding in integers into a generator instead of splitting it?

from jax.random import PRNGKey, split, fold_in
from jax import enable_custom_prng

with enable_custom_prng():
    g = PRNGKey(123)
    h = PRNGKey(123)

for rng in split(g, 10):  # A
    print(rng)

print()

for i in range(10):  # B
    print(fold_in(h, i))

[froystig]

We had precisely the example B in mind for fold_in usage. I would go with that!

Is there any downside to folding in integers into a generator instead of splitting it?

If anything, there are possible advantages for working memory.

[eadadi]

Hey, adding a question -
Is there any problem with folding a folded_key?
Below I present an example that shows sensible output, but I think I read somewhere folded keys are not intended to be relied for making new keys.

I imagine something like:

for i in range(3):
  sk = random.fold_in(k, i)
    for j in range(3):
      sksk = random.fold_in(sk, j)
        print(f"k={k}, i={i}, sk={sk}, j={j}, sksk={sksk}")

Which gives sensible output:

k=[0 0], i=0, sk=[1797259609 2579123966], j=0, sksk=[4165894930  804218099]
k=[0 0], i=0, sk=[1797259609 2579123966], j=1, sksk=[1353695780 2116000888]
k=[0 0], i=0, sk=[1797259609 2579123966], j=2, sksk=[2839387376 2467677468]
k=[0 0], i=1, sk=[ 928981903 3453687069], j=0, sksk=[346279018 360566543]
k=[0 0], i=1, sk=[ 928981903 3453687069], j=1, sksk=[3968330031 3923691647]
k=[0 0], i=1, sk=[ 928981903 3453687069], j=2, sksk=[3152301319  792466127]
k=[0 0], i=2, sk=[4146024105 2718843009], j=0, sksk=[2799984767 1105366846]
k=[0 0], i=2, sk=[4146024105 2718843009], j=1, sksk=[3777617834  145086855]
k=[0 0], i=2, sk=[4146024105 2718843009], j=2, sksk=[ 915694800 1641710144]

[NeilGirdhar]

My guess is that by induction, there should be no problem. After all, each sk is a reasonably key, so it should be a reasonable start to a key sequence (sk, j).

Also, I think Flax's RngKeySequence used to do something like this.