Parallel Random Number Generation

[Seb-Lorek]

I have a general question regarding random number generation in JAX. Does pseudo random number generation in JAX automatically produce repeatable pseudo-random numbers across multiple processes (local or distributed) such that the streams are independent. In numpy this can be guaranteed (with very very high prob) by spawning. Im interest into this due to using JAX pseudo random number generation in a simulation study. Let's say my run_sim() function looks something like this (no parallelization so far, later on i want to parallelize it to run on multiple CPU cores)

# Function to execute the simulation study 
def sim_fun(n_sim):
    # Create a grid of different model specifications and data generating mechanisms (f.e. n_obs, data relationship etc.)
    var_grid = create_var_grid(n_sim)
    # I need two ints as seeds, one seed for the variational inference algorithm and one seed for the data generation 
    split = jnp.split(jnp.arange(n_sim*2*len(var_grid)), len(var_grid))
    seeds = [jnp.split(subsplit, n_sim) for subsplit in split]
    sim_results = []
    true_coefs = []

    for j, config in enumerate(var_grid):
        store_result = []
        store_coef = []
        # I want to parallelize this for loop to use multiple CPU cores 
        for i in range(n_sim):
            # Run the simulation once for a configuration and pass the two ints to generate two new keys 
            result, true_coef = do_one(config, seeds[j][i])
            store_result.append(result)
            store_coef.append(true_coef)
        true_coefs.append(store_coef)
        sim_results.append(store_result)

    return sim_results, true_coefs

Are the random number streams independent or are they overlapping with such a set up? Probably JAX takes care of this, however i couldn't find something about this in the docs. And is this a good practice, i.e. to define an array of seeds?

[froystig]

If you start with a JAX RNG key k, you can use jax.random.split or jax.random.fold_in to derive keys from k that are intended to produce (pseudo) independent streams, i.e. ideally uncorrelated pseudorandom samples (directly or indirectly via further key derivation).

[Seb-Lorek]

I see so this results in independent streams

key = jax.random.PRNGKey(42)
# generate 250 new keys for each iteration in the simulation
key, *subkeys = jax.random.split(key, 251)

Each of the new keys generates ideally uncorrelated (pseudo) random samples. However essentially in the example above inside the do_one()function I would gerate two keys from the integers (jnp.arange())

# generate two new keys in each iteration with a new seed
key_inference = jax.random.PRNGKey(0)
key_data_sim = jax.random.PRNGKey(1)

and split them downwards from there for new random numbers. Does this also result in uncorrelated streams or should i rather generate one key and generate childs (subkeys) for each iteration of the simulation and then grandchildren ... downwards .

[froystig]

Any key derivation (i.e. split) is intended to produce uncorrelated streams from its subkeys. That includes splits from splits.

For your use case, fold_in might make more sense. Specifically, you could considering folding in the current iteration count.

[Seb-Lorek]

Nice, thanks!

[jecampagne]

Hello,
I'm looking for a description of what is the underlaying pseudo random number generator used by JAX in CPU and GPU. Does it is some mersenne twister algorithm or a falvor of RANLUX?
Thanks

[jakevdp]

Hi - you can read about JAX's PRNG design here: https://jax.readthedocs.io/en/latest/jep/263-prng.html

[jecampagne]

Thanks @jakevdp