atic"
prompt = [prompt] * jax.device_count()
prompt_ids = pipeline.prepare_inputs(prompt)
prompt_ids.shape
# (8, 77) Model parameters and inputs have to be replicated across the 8 parallel devices. The parameters dictionary is replicated with flax.jax_utils.replicate which traverses the dictionary and changes the shape of the weights so they are repeated 8 times. Arrays are replicated using shard.   Copied # parameters
p_params = replicate(para