This is a great idea. I agree having both gather and allgather for variable size is very useful. Also, seeing as the displacement is an implementation detail, the underlying C++ implementation can take care of sharing each process’ contribution, allocating the required memory, and returning the list of tensors.
I created https://github.com/pytorch/pytorch/issues/23299 8 to track the feature.