DOC: add note about same sized tensors to dist.gather() (#128676)
Fixes #103305
Pull Request resolved: https://github.com/pytorch/pytorch/pull/128676
Approved by: https://github.com/wconstab
diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py
index bd81fd6..d44c373 100644
--- a/torch/distributed/distributed_c10d.py
+++ b/torch/distributed/distributed_c10d.py
@@ -3041,11 +3041,12 @@
"""
Gathers tensors from the whole group in a list.
- Complex tensors are supported.
+ Complex and uneven sized tensors are supported.
Args:
tensor_list (list[Tensor]): Output list. It should contain
correctly-sized tensors to be used for output of the collective.
+ Uneven sized tensors are supported.
tensor (Tensor): Tensor to be broadcast from current process.
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used.
@@ -3118,6 +3119,8 @@
"""
Gather tensors from all ranks and put them in a single output tensor.
+ This function requires all tensors to be the same size on each process.
+
Args:
output_tensor (Tensor): Output tensor to accommodate tensor elements
from all ranks. It must be correctly sized to have one of the
@@ -3341,11 +3344,13 @@
"""
Gathers a list of tensors in a single process.
+ This function requires all tensors to be the same size on each process.
+
Args:
tensor (Tensor): Input tensor.
- gather_list (list[Tensor], optional): List of appropriately-sized
- tensors to use for gathered data (default is None, must be specified
- on the destination rank)
+ gather_list (list[Tensor], optional): List of appropriately,
+ same-sized tensors to use for gathered data
+ (default is None, must be specified on the destination rank)
dst (int, optional): Destination rank on global process group (regardless of ``group`` argument). (default is 0)
group (ProcessGroup, optional): The process group to work on. If None,
the default process group will be used.
