eplication and parallelization
	
Model parameters and inputs have to be replicated across the 8 parallel devices we have. The parameters dictionary is replicated using 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
p_params = replicate(params)


	Copied
prompt_ids = shard(prompt_ids)
prompt_ids.shape


	Copied
(8, 1, 77)
Tha