Update all_gather docs (#19469)

Author: awaelchli

docs/source-fabric/advanced/distributed_communication.rst

@@ -236,6 +236,11 @@ Full example:
     result = fabric.all_gather(data)
     print("Result of all-gather:", result)  # tensor([ 0, 10, 20, 30])
 
+.. warning::
+
+    For the special case where ``world_size`` is 1, no additional dimension is added to the tensor(s). This inconsistency
+    is kept for backward compatibility and you may need to handle this special case in your code to make it agnostic.
+
 
 ----
 

src/lightning/fabric/fabric.py

@@ -583,7 +583,8 @@ def all_gather(
 
         Return:
             A tensor of shape (world_size, batch, ...), or if the input was a collection
-            the output will also be a collection with tensors of this shape.
+            the output will also be a collection with tensors of this shape. For the special case where
+            world_size is 1, no additional dimension is added to the tensor(s).
 
         """
         self._validate_launched()

src/lightning/pytorch/core/module.py

@@ -668,7 +668,8 @@ def all_gather(
 
         Return:
             A tensor of shape (world_size, batch, ...), or if the input was a collection
-            the output will also be a collection with tensors of this shape.
+            the output will also be a collection with tensors of this shape. For the special case where
+            world_size is 1, no additional dimension is added to the tensor(s).
 
         """
         group = group if group is not None else torch.distributed.group.WORLD