Feel free to use different seeds when exploring results later in the guide.   Copied def create_key(seed=0):
    return jax.random.PRNGKey(seed) The helper function, or rng, is split 8 times so each device receives a different generator and generates a different image.   Copied rng = create_key(0)
rng = jax.random.split(rng, jax.device_count()) To take advantage of JAX’s optimized speed on a TPU, pass jit=True to the pipeline to compile the JAX