Allow passing a jax.random.PRNGKey to class EngineBuilder

[Seb-Lorek]

Hey liesel devs. Currently I'm pondering on how to safely use parallel pseudo-random number generation in JAX for a simulation study. In a simulation study in general it is important to have 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. In JAX we can either use jax.random.split or jax.random.fold_in for safe production of (pseudo) independent streams (directly or indirectly via further key derivation). However both involve passing a jax.random.PRNGKey and not the seed. Creating a iterable/list of integers and iterating over them is not safe, corresponding to the following

for i in range(n):
    rng = jax.random.PRNGKey(i)

See #18211 and here for some discussions.

In order to reliably produce independent streams within a simulation study, it would be important to have the option of directly passing a jax.random.PRNGKey to the EngineBuilder. This should be quite easy to implement as far as I can see.

[wiep]

Passing a KeyArray was possible at some point. I do not know why we removed it. However, it should be pretty straightforward implement it in a PR. Please feel encourage to give it a try.

In the constructor of EngineBuilder, you want to check the type of the seed argument:

• If it is a KeyArray / PRNGKey use it directly in jax.split
• if it is a valid argument to PRNGKey create a new key and then split it
• Otherwise: raise error.

liesel/liesel/goose/builder.py

Lines 132 to 133 in 1ee02de

	def __init__(self, seed: int, num_chains: int):
	keys = jax.random.split(jax.random.PRNGKey(seed), 3)

[Seb-Lorek]

Hey sure I'll give it a try later.

[jobrachem]

Thanks for your contribution! Since the request has since been implemented in #157, I'm closing this thread.