pmap: data size non-divisible by device_count on data parallel cases

[luweizheng]

Hi all,

I am using pmap and it is very powerful to parallelize data onto multiple devices.
To use pmap, we should divide data_size by jax.local_device_count(). In this case, data is split onto different devices.
However, it is very common that data_size may not be divisible by local_device_count.
I want to keep the remaining rather than drop them. So the common practice is to fill some dummy elements into the data and make data_size divisible. After pmap, we should mask out the dummy elements.
My question is: Are there any examples to do this fill and masked out thing? As the result from pmap is a ShardedDeviceArray, can I just mask out the dummy elements on the last device?

Thanks.