Repository: pytorch/ignite

Issue:
Improve XLA tensor logging
## 🚀 Feature

Following https://github.com/pytorch/xla/issues/2225, we need to improve our code of experiment tracking loggers when run on XLA devices.



Discussion:

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
@@ -119,10 +119,13 @@ def _all_collective_op(
             tensor = self._encode_str(tensor, device)
 
         out_dtype = None
+        tensor_device = None
 
         # check if the tensor is at specified device
         if tensor.device != device:
+            tensor_device = tensor.device
             tensor = tensor.to(device)
+
         if self._collective_op_dtype is not None:
             # cast to _collective_op_dtype if current type is not floatX
             if tensor.dtype not in (torch.float32, torch.float64):
@@ -131,8 +134,12 @@ def _all_collective_op(
 
         tensor = fn(tensor, *args, **kwargs)
 
-        if out_dtype is not None:
+        if out_dtype is not None and tensor_device is not None:
+            tensor = tensor.to(dtype=out_dtype, device=tensor_device)
+        elif out_dtype is not None:
             tensor = tensor.to(dtype=out_dtype)
+        elif tensor_device is not None:
+            tensor = tensor.to(device=tensor_device)
 
         if tensor_to_number and tensor.numel() == 1:
             return tensor.item()
diff --git a/ignite/metrics/precision.py b/ignite/metrics/precision.py
--- a/ignite/metrics/precision.py
+++ b/ignite/metrics/precision.py
@@ -44,7 +44,7 @@ def reset(self) -> None:
         self._positives = torch.tensor([], dtype=dtype) if (self._is_multilabel and not self._average) else 0
         super(_BasePrecisionRecall, self).reset()
 
-    def compute(self) -> torch.Tensor:
+    def compute(self) -> Union[torch.Tensor, float]:
         if not (isinstance(self._positives, torch.Tensor) or self._positives > 0):
             raise NotComputableError(
                 "{} must have at least one example before" " it can be computed.".format(self.__class__.__name__)


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
@@ -74,6 +74,10 @@ def _test_distrib_all_reduce(device):
         with pytest.raises(ValueError, match=r"Unsupported reduction operation"):
             idist.all_reduce(10, op="ABC")
 
+        t = torch.tensor([0, 1, 2])
+        res = idist.all_reduce(t)
+        assert res.device == t.device, "{} vs {}".format(res.device, t.device)
+
 
 def _test_distrib_all_gather(device):
 
diff --git a/tests/ignite/metrics/test_precision.py b/tests/ignite/metrics/test_precision.py
--- a/tests/ignite/metrics/test_precision.py
+++ b/tests/ignite/metrics/test_precision.py
@@ -739,7 +739,7 @@ def update(engine, i):
 
         engine = Engine(update)
 
-        pr = Precision(average=average, device=device)
+        pr = Precision(average=average)
         pr.attach(engine, "pr")
 
         data = list(range(n_iters))
@@ -748,6 +748,7 @@ def update(engine, i):
         assert "pr" in engine.state.metrics
         res = engine.state.metrics["pr"]
         if isinstance(res, torch.Tensor):
+            assert res.device.type == "cpu"
             res = res.cpu().numpy()
 
         true_res = precision_score(
@@ -787,7 +788,7 @@ def update(engine, i):
 
         engine = Engine(update)
 
-        pr = Precision(average=average, is_multilabel=True, device=device)
+        pr = Precision(average=average, is_multilabel=True)
         pr.attach(engine, "pr")
 
         data = list(range(n_iters))
@@ -821,7 +822,7 @@ def update(engine, i):
             match="Precision/Recall metrics do not work in distributed setting when "
             "average=False and is_multilabel=True",
         ):
-            pr = Precision(average=False, is_multilabel=True, device=device)
+            pr = Precision(average=False, is_multilabel=True)
 
         y_pred = torch.randint(0, 2, size=(4, 3, 6, 8))
         y = torch.randint(0, 2, size=(4, 3, 6, 8)).long()
diff --git a/tests/ignite/metrics/test_recall.py b/tests/ignite/metrics/test_recall.py
--- a/tests/ignite/metrics/test_recall.py
+++ b/tests/ignite/metrics/test_recall.py
@@ -739,7 +739,7 @@ def update(engine, i):
 
         engine = Engine(update)
 
-        re = Recall(average=average, device=device)
+        re = Recall(average=average)
         re.attach(engine, "re")
 
         data = list(range(n_iters))
@@ -748,6 +748,7 @@ def update(engine, i):
         assert "re" in engine.state.metrics
         res = engine.state.metrics["re"]
         if isinstance(res, torch.Tensor):
+            assert res.device.type == "cpu"
             res = res.cpu().numpy()
 
         true_res = recall_score(
@@ -787,7 +788,7 @@ def update(engine, i):
 
         engine = Engine(update)
 
-        re = Recall(average=average, is_multilabel=True, device=device)
+        re = Recall(average=average, is_multilabel=True)
         re.attach(engine, "re")
 
         data = list(range(n_iters))
@@ -821,7 +822,7 @@ def update(engine, i):
             match="Precision/Recall metrics do not work in distributed setting when "
             "average=False and is_multilabel=True",
         ):
-            re = Recall(average=False, is_multilabel=True, device=device)
+            re = Recall(average=False, is_multilabel=True)
 
         y_pred = torch.randint(0, 2, size=(4, 3, 6, 8))
         y = torch.randint(0, 2, size=(4, 3, 6, 8)).long()