p, the prepared function p_generate will conceptually do the following:
Invoke a copy of the underlying function pipeline._generate in each device.
Send each device a different portion of the input arguments. That’s what sharding is used for. In our case, prompt_ids has shape (8, 1, 77, 768). This array will be split in 8 and each copy of _generate will receive an input with shape (1, 77, 768).
We can code _generate completely ignoring the fact