p(pipeline._generate) After calling pmap, the prepared function p_generate will: Make a copy of the underlying function, pipeline._generate, on each device. Send each device a different portion of the input arguments (this is why it’s necessary to call the shard function). In this case, prompt_ids has shape (8, 1, 77, 768) so the array is split into 8 and each copy of _generate receives an input with shape (1, 77, 768). The most important thing