ent seeds when exploring results later in the notebook.


	Copied
def create_key(seed=0):
    return jax.random.PRNGKey(seed)
We obtain a rng and then “split” it 8 times so each device receives a different generator. Therefore, each device will create a different image, and the full process is reproducible.


	Copied
rng = create_key(0)
rng = jax.random.split(rng, jax.device_count())
JAX code can be compiled to an efficient representation that 