Repository: pytorch/ignite

Issue:
Improve typing of distributed.comp_modules.utils.all_gather
## 🚀 Feature

As a followup to the issue #1344 and comment https://github.com/pytorch/ignite/pull/1355#discussion_r499854850 it is needed to rework the typing of the `utils.all_gather` and the different realizations serial, native, hvd, xla.



Discussion:
@vfdev-5 I took a look on the code again and your suggested typing, is not feasible.
```
def all_gather(tensor: Union[torch.Tensor, Number, str]) -> Union[torch.Tensor, List[Number], List[str]]:
```
Looking at the definition of `all_gather` inside the `ComputationModel` class
https://github.com/pytorch/ignite/blob/380dd3b0de3c8b4459b2ef221fcefb29f42f9458/ignite/distributed/comp_models/base.py#L163-L167
there is a reference to the function `_collective_op` and there is no possibility to get a result of type `List[Number]`
https://github.com/pytorch/ignite/blob/380dd3b0de3c8b4459b2ef221fcefb29f42f9458/ignite/distributed/comp_models/base.py#L137-L155
except I change the return in line 152 to
```
if tensor_to_number and tensor.numel() == 1:
    return cast(List[Number], [tensor.item()])
```
If it is ok for you, then I would make the change and adjust all the typings and create the PR.
@gruebel you are right, `_collective_op` implementation does not handle this case properly. Let me detail it, in case of all gather, collective op does the following for the input as number:
```
Number -> Tensor[1] -> Tensor[N] returns Tensor[N]
```
comparing with in case of all reduce, collective op does:
```
Number -> Tensor[1] -> Tensor[1] -> Number returns Number
```

Probably, we can add the following case
```python
if tensor_to_number:
    if tensor.numel() == 1:
        return cast(Number, tensor.item())
    else:
        return case(List[Number], tensor.tolist())
```

I think "return" type annotations of `_collective_op` is not complete, I'd say it should be 
```python
    def _collective_op(...) -> Union[torch.Tensor, Number, List[Number], List[str]]:
```
and cover all reduce and all gather cases.

ok, that makes sense. I also saw the reference to `~Tensor.tolist` there. I will give it try 💪 

Resolving patch:
diff --git a/ignite/distributed/comp_models/base.py b/ignite/distributed/comp_models/base.py
--- a/ignite/distributed/comp_models/base.py
+++ b/ignite/distributed/comp_models/base.py
@@ -136,7 +136,7 @@ def _apply_op(
 
     def _collective_op(
         self, tensor: Union[torch.Tensor, Number, str], fn: Callable, *args: Any, **kwargs: Any
-    ) -> Union[torch.Tensor, Number, List[str]]:
+    ) -> Union[torch.Tensor, Number, List[Number], List[str]]:
         tensor_to_number = tensor_to_str = False
         device = self.device()
         if isinstance(tensor, Number):
@@ -148,8 +148,11 @@ def _collective_op(
 
         tensor = self._apply_op(tensor, device, fn, *args, **kwargs)
 
-        if tensor_to_number and tensor.numel() == 1:
-            return cast(Number, tensor.item())
+        if tensor_to_number:
+            if tensor.numel() == 1:
+                return cast(Number, tensor.item())
+            else:
+                return tensor.tolist()
         elif tensor_to_str:
             return self._decode_str(tensor)
         return tensor
@@ -160,7 +163,9 @@ def all_reduce(self, tensor: Union[torch.Tensor, Number], op: str = "sum") -> Un
 
         return cast(Union[torch.Tensor, Number], self._collective_op(tensor, self._do_all_reduce, op))
 
-    def all_gather(self, tensor: Union[torch.Tensor, Number, str]) -> Union[torch.Tensor, Number, List[str]]:
+    def all_gather(
+        self, tensor: Union[torch.Tensor, Number, str]
+    ) -> Union[torch.Tensor, Number, List[Number], List[str]]:
         if not isinstance(tensor, (torch.Tensor, Number, str)):
             raise TypeError("Unhandled input type {}".format(type(tensor)))
 
@@ -271,8 +276,12 @@ def spawn(*args: Any, **kwargs: Any) -> None:
     def all_reduce(self, tensor: Union[torch.Tensor, Number], op: str = "sum") -> Union[torch.Tensor, Number]:
         return tensor
 
-    def all_gather(self, tensor: Union[torch.Tensor, Number]) -> Union[torch.Tensor, Number]:  # type: ignore
-        return tensor
+    def all_gather(
+        self, tensor: Union[torch.Tensor, Number, str]
+    ) -> Union[torch.Tensor, Number, List[Number], List[str]]:
+        if isinstance(tensor, torch.Tensor):
+            return tensor
+        return cast(Union[List[Number], List[str]], [tensor])
 
     def broadcast(self, tensor: Union[torch.Tensor, Number, str], src: int = 0) -> Union[torch.Tensor, Number, str]:
         return tensor
diff --git a/ignite/distributed/utils.py b/ignite/distributed/utils.py
--- a/ignite/distributed/utils.py
+++ b/ignite/distributed/utils.py
@@ -334,7 +334,7 @@ def all_reduce(tensor: Union[torch.Tensor, Number], op: str = "SUM") -> Union[to
     return _model.all_reduce(tensor, op)
 
 
-def all_gather(tensor: Union[torch.Tensor, Number, str]) -> Union[torch.Tensor, Number, List[str]]:
+def all_gather(tensor: Union[torch.Tensor, Number, str]) -> Union[torch.Tensor, Number, List[Number], List[str]]:
     """Helper method to perform all gather operation.
 
     Args:
@@ -349,7 +349,7 @@ def all_gather(tensor: Union[torch.Tensor, Number, str]) -> Union[torch.Tensor,
     if _need_to_sync and isinstance(_model, _SerialModel):
         sync(temporary=True)
 
-    return _model.all_gather(tensor)  # type: ignore[arg-type]
+    return _model.all_gather(tensor)
 
 
 def broadcast(tensor: Union[torch.Tensor, Number, str], src: int = 0) -> Union[torch.Tensor, Number, str]:


Test patch:
diff --git a/tests/ignite/distributed/utils/__init__.py b/tests/ignite/distributed/utils/__init__.py
--- a/tests/ignite/distributed/utils/__init__.py
+++ b/tests/ignite/distributed/utils/__init__.py
@@ -83,7 +83,7 @@ def _test_distrib_all_reduce(device):
 
 def _test_distrib_all_gather(device):
 
-    res = idist.all_gather(10)
+    res = torch.tensor(idist.all_gather(10), device=device)
     true_res = torch.tensor([10,] * idist.get_world_size(), device=device)
     assert (res == true_res).all()
 
diff --git a/tests/ignite/distributed/utils/test_serial.py b/tests/ignite/distributed/utils/test_serial.py
--- a/tests/ignite/distributed/utils/test_serial.py
+++ b/tests/ignite/distributed/utils/test_serial.py
@@ -54,4 +54,5 @@ def test_idist_all_reduce_no_dist():
 
 
 def test_idist_all_gather_no_dist():
-    assert idist.all_gather(10) == 10
+    assert idist.all_gather(10) == [10]
+    assert (idist.all_gather(torch.tensor(10)) == torch.tensor(10)).all()