vmap + PRNG Keys

You want jax.vmap to apply a jitted function independently to a batch of 32 inputs, where each call needs its own fresh random key for dropout. What's the correct way to supply the randomness?

Call jax.random.PRNGKey(i) inside the function body using the vmap batch index

Loop in Python, calling jax.random.fold_in(key, i) for each of the 32 elements sequentially

Pass one shared key into every vmapped call, using it as-is inside the function

Split a top-level key with jax.random.split(key, 32) and vmap over the resulting keys array with in_axes=0

Reuse the same key for every element and vmap over the dropout rate instead