shard.


	Copied
p_params = replicate(params)


	Copied
prompt_ids = shard(prompt_ids)
prompt_ids.shape


	Copied
(8, 1, 77)
That shape means that each one of the 8 devices will receive as an input a jnp array with shape (1, 77). 1 is therefore the batch size per device. In TPUs with sufficient memory, it could be larger than 1 if we wanted to generate multiple images (per chip) at once.
We are almost ready to generate images! We just need to c