update xla tests

tests/ignite/metrics/test_accumulation.py

@@ -370,7 +370,10 @@ def _geom_mean(y_true):
 
 def _test_distrib_accumulator_device(device):
 
-    for metric_device in [torch.device("cpu"), idist.device()]:
+    metric_devices = [torch.device("cpu")]
+    if device.type != "xla":
+        metric_devices.append(device)
+    for metric_device in metric_devices:
 
         m = VariableAccumulation(lambda a, x: x, device=metric_device)
         assert m._device == metric_device
@@ -384,11 +387,6 @@ def _test_distrib_accumulator_device(device):
         )
 
 
-def _test_creating_on_xla_fails(device):
-    with pytest.raises(ValueError, match=r"Cannot create metric on an XLA device. Use device='cpu' instead."):
-        VariableAccumulation(lambda a, x: x, device=device)
-
-
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
@@ -462,7 +460,7 @@ def test_distrib_single_device_xla():
     _test_distrib_average(device)
     _test_distrib_geom_average(device)
     _test_distrib_integration(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 def _test_distrib_xla_nprocs(index):
@@ -471,7 +469,7 @@ def _test_distrib_xla_nprocs(index):
     _test_distrib_average(device)
     _test_distrib_geom_average(device)
     _test_distrib_integration(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu

tests/ignite/metrics/test_accuracy.py

@@ -813,7 +813,10 @@ def update(engine, i):
 
 def _test_distrib_accumulator_device(device):
 
-    for metric_device in [torch.device("cpu"), idist.device()]:
+    metric_devices = [torch.device("cpu")]
+    if device.type != "xla":
+        metric_devices.append(device)
+    for metric_device in metric_devices:
 
         acc = Accuracy(device=metric_device)
         assert acc._device == metric_device
@@ -830,11 +833,6 @@ def _test_distrib_accumulator_device(device):
         )
 
 
-def _test_creating_on_xla_fails(device):
-    with pytest.raises(ValueError, match=r"Cannot create metric on an XLA device. Use device='cpu' instead."):
-        Accuracy(device=device)
-
-
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
@@ -901,15 +899,15 @@ def test_distrib_single_device_xla():
     _test_distrib_multilabel_input_NHW(device)
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 def _test_distrib_xla_nprocs(index):
     device = idist.device()
     _test_distrib_multilabel_input_NHW(device)
     _test_distrib_integration_multiclass(device)
     _test_distrib_integration_multilabel(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu

tests/ignite/metrics/test_confusion_matrix.py

@@ -611,12 +611,16 @@ def _test(metric_device):
         assert np.all(true_res == res)
 
     _test("cpu")
-    _test(idist.device())
+    if device.type != "xla":
+        _test(idist.device())
 
 
 def _test_distrib_accumulator_device(device):
 
-    for metric_device in [torch.device("cpu"), idist.device()]:
+    metric_devices = [torch.device("cpu")]
+    if device.type != "xla":
+        metric_devices.append(device)
+    for metric_device in metric_devices:
 
         cm = ConfusionMatrix(num_classes=3, device=metric_device)
         assert cm._device == metric_device
@@ -638,7 +642,7 @@ def _test_distrib_accumulator_device(device):
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
 def test_distrib_gpu(local_rank, distributed_context_single_node_nccl):
 
-    device = "cuda:{}".format(local_rank)
+    device = torch.device("cuda:{}".format(local_rank))
     _test_distrib_multiclass_images(device)
     _test_distrib_accumulator_device(device)
 
@@ -647,7 +651,7 @@ def test_distrib_gpu(local_rank, distributed_context_single_node_nccl):
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 def test_distrib_cpu(distributed_context_single_node_gloo):
 
-    device = "cpu"
+    device = torch.device("cpu")
     _test_distrib_multiclass_images(device)
     _test_distrib_accumulator_device(device)
 
@@ -657,7 +661,7 @@ def test_distrib_cpu(distributed_context_single_node_gloo):
 @pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
 def test_distrib_hvd(gloo_hvd_executor):
 
-    device = "cpu" if not torch.cuda.is_available() else "cuda"
+    device = torch.device("cpu" if not torch.cuda.is_available() else "cuda")
     nproc = 4 if not torch.cuda.is_available() else torch.cuda.device_count()
 
     gloo_hvd_executor(_test_distrib_multiclass_images, (device,), np=nproc, do_init=True)
@@ -668,7 +672,7 @@ def test_distrib_hvd(gloo_hvd_executor):
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif("MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
 def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
-    device = "cpu"
+    device = torch.device("cpu")
     _test_distrib_multiclass_images(device)
     _test_distrib_accumulator_device(device)
 
@@ -677,7 +681,7 @@ def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif("GPU_MULTINODE_DISTRIB" not in os.environ, reason="Skip if not multi-node distributed")
 def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
-    device = "cuda:{}".format(distributed_context_multi_node_nccl["local_rank"])
+    device = torch.device("cuda:{}".format(distributed_context_multi_node_nccl["local_rank"]))
     _test_distrib_multiclass_images(device)
     _test_distrib_accumulator_device(device)
 

tests/ignite/metrics/test_loss.py

@@ -111,9 +111,12 @@ def _test(metric_device):
         _test(device)
 
 
-def _test_distrib_sum_device(device):
+def _test_distrib_accumulator_device(device):
 
-    for metric_device in [torch.device("cpu"), idist.device()]:
+    metric_devices = [torch.device("cpu")]
+    if device.type != "xla":
+        metric_devices.append(device)
+    for metric_device in metric_devices:
         loss = Loss(nll_loss, device=metric_device)
         assert loss._device == metric_device
         assert loss._sum.device == metric_device, "{}:{} vs {}:{}".format(
@@ -139,19 +142,14 @@ def test_sum_detached():
     assert not loss._sum.requires_grad
 
 
-def _test_creating_on_xla_fails(device):
-    with pytest.raises(ValueError, match=r"Cannot create metric on an XLA device. Use device='cpu' instead."):
-        Loss(nll_loss, device=device)
-
-
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
 def test_distrib_gpu(local_rank, distributed_context_single_node_nccl):
 
     device = torch.device("cuda:{}".format(local_rank))
     _test_distrib_compute_on_criterion(device)
-    _test_distrib_sum_device(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.distributed
@@ -160,7 +158,7 @@ def test_distrib_cpu(distributed_context_single_node_gloo):
 
     device = torch.device("cpu")
     _test_distrib_compute_on_criterion(device)
-    _test_distrib_sum_device(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.distributed
@@ -172,7 +170,7 @@ def test_distrib_hvd(gloo_hvd_executor):
     nproc = 4 if not torch.cuda.is_available() else torch.cuda.device_count()
 
     gloo_hvd_executor(_test_distrib_compute_on_criterion, (device,), np=nproc, do_init=True)
-    gloo_hvd_executor(_test_distrib_sum_device, (device,), np=nproc, do_init=True)
+    gloo_hvd_executor(_test_distrib_accumulator_device, (device,), np=nproc, do_init=True)
 
 
 @pytest.mark.multinode_distributed
@@ -181,7 +179,7 @@ def test_distrib_hvd(gloo_hvd_executor):
 def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
     device = torch.device("cpu")
     _test_distrib_compute_on_criterion(device)
-    _test_distrib_sum_device(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.multinode_distributed
@@ -190,7 +188,7 @@ def test_multinode_distrib_cpu(distributed_context_multi_node_gloo):
 def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
     device = torch.device("cuda:{}".format(distributed_context_multi_node_nccl["local_rank"]))
     _test_distrib_compute_on_criterion(device)
-    _test_distrib_sum_device(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu
@@ -199,13 +197,13 @@ def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
 def test_distrib_single_device_xla():
     device = idist.device()
     _test_distrib_compute_on_criterion(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 def _test_distrib_xla_nprocs(index):
     device = idist.device()
     _test_distrib_compute_on_criterion(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu

tests/ignite/metrics/test_mean_absolute_error.py

@@ -72,7 +72,10 @@ def _test(metric_device):
 
 def _test_distrib_accumulator_device(device):
 
-    for metric_device in [torch.device("cpu"), idist.device()]:
+    metric_devices = [torch.device("cpu")]
+    if device.type != "xla":
+        metric_devices.append(device)
+    for metric_device in metric_devices:
         mae = MeanAbsoluteError(device=metric_device)
         assert mae._device == metric_device
         assert mae._sum_of_absolute_errors.device == metric_device, "{}:{} vs {}:{}".format(
@@ -103,11 +106,6 @@ def test_accumulator_detached():
     assert not mae._sum_of_absolute_errors.requires_grad
 
 
-def _test_creating_on_xla_fails(device):
-    with pytest.raises(ValueError, match=r"Cannot create metric on an XLA device. Use device='cpu' instead."):
-        MeanAbsoluteError(device=device)
-
-
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
@@ -161,13 +159,13 @@ def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
 def test_distrib_single_device_xla():
     device = idist.device()
     _test_distrib_integration(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 def _test_distrib_xla_nprocs(index):
     device = idist.device()
     _test_distrib_integration(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu

tests/ignite/metrics/test_mean_pairwise_distance.py

@@ -85,7 +85,10 @@ def _test(metric_device):
 
 def _test_distrib_accumulator_device(device):
 
-    for metric_device in [torch.device("cpu"), idist.device()]:
+    metric_devices = [torch.device("cpu")]
+    if device.type != "xla":
+        metric_devices.append(device)
+    for metric_device in metric_devices:
 
         mpd = MeanPairwiseDistance(device=metric_device)
         assert mpd._device == metric_device
@@ -112,11 +115,6 @@ def test_accumulator_detached():
     assert not mpd._sum_of_distances.requires_grad
 
 
-def _test_creating_on_xla_fails(device):
-    with pytest.raises(ValueError, match=r"Cannot create metric on an XLA device. Use device='cpu' instead."):
-        MeanPairwiseDistance(device=device)
-
-
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
@@ -170,13 +168,13 @@ def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
 def test_distrib_single_device_xla():
     device = idist.device()
     _test_distrib_integration(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 def _test_distrib_xla_nprocs(index):
     device = idist.device()
     _test_distrib_integration(device)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu

tests/ignite/metrics/test_mean_squared_error.py

@@ -72,7 +72,10 @@ def _test(metric_device):
 
 def _test_distrib_accumulator_device(device):
 
-    for metric_device in [torch.device("cpu"), idist.device()]:
+    metric_devices = [torch.device("cpu")]
+    if device.type != "xla":
+        metric_devices.append(device)
+    for metric_device in metric_devices:
 
         device = torch.device(device)
         mse = MeanSquaredError(device=metric_device)
@@ -105,11 +108,6 @@ def test_accumulator_detached():
     assert not mse._sum_of_squared_errors.requires_grad
 
 
-def _test_creating_on_xla_fails(device):
-    with pytest.raises(ValueError, match=r"Cannot create metric on an XLA device. Use device='cpu' instead."):
-        MeanSquaredError(device=device)
-
-
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
@@ -164,13 +162,13 @@ def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
 def test_distrib_single_device_xla():
     device = idist.device()
     _test_distrib_integration(device, tol=1e-4)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 def _test_distrib_xla_nprocs(index):
     device = idist.device()
     _test_distrib_integration(device, tol=1e-4)
-    _test_creating_on_xla_fails(device)
+    _test_distrib_accumulator_device(device)
 
 
 @pytest.mark.tpu

tests/ignite/metrics/test_metric.py

@@ -568,13 +568,19 @@ def update(self, output):
             self.a += 10.0
             self.b -= 5.0
 
-    m = DummyMetric(device=device)
+    metric_device = device if torch.device(device).type != "xla" else "cpu"
+    m = DummyMetric(device=metric_device)
     m.update(None)
     m.compute()
     # check if can call compute multiple times without all reduce invocation
     m.compute()
 
 
+def _test_creating_on_xla_fails(device):
+    with pytest.raises(ValueError, match=r"Cannot create metric on an XLA device. Use device='cpu' instead."):
+        DummyMetric2(device=device)
+
+
 @pytest.mark.distributed
 @pytest.mark.skipif(not idist.has_native_dist_support, reason="Skip if no native dist support")
 @pytest.mark.skipif(torch.cuda.device_count() < 1, reason="Skip if no GPU")
@@ -625,11 +631,13 @@ def test_multinode_distrib_gpu(distributed_context_multi_node_nccl):
 def test_distrib_single_device_xla():
     device = idist.device()
     _test_distrib_sync_all_reduce_decorator(device)
+    _test_creating_on_xla_fails(device)
 
 
 def _test_distrib_xla_nprocs(index):
     device = idist.device()
     _test_distrib_sync_all_reduce_decorator(device)
+    _test_creating_on_xla_fails(device)
 
 
 @pytest.mark.tpu