# (8, 1, 77) This shape means each one of the 8 devices receives as an input a jnp array with shape (1, 77), where 1 is the batch size per device. On TPUs with sufficient memory, you could have a batch size larger than 1 if you want to generate multiple images (per chip) at once. Next, create a random number generator to pass to the generation function. This is standard procedure in Flax, which is very serious and opinionated about random numbers. All functions that deal with random numbers are expected to receive a generator to ensure reproducibility, even when you’re training across multiple distributed devices. The helper function below uses a seed to initialize a random number generator. As long as you use the same seed, you’ll get the exact same results. Feel free to use different seeds when exploring results later in the guide.   Copied def create_key(seed=0):