Repository: pytorch/ignite

Issue:
Improve `ignite.contirb.metrics` implementation 
For these metrics in `ignite.contrib.metrics` :
- [Average Precision](https://github.com/pytorch/ignite/blob/master/ignite/contrib/metrics/average_precision.py)
- [Precision Recall Curve](https://github.com/pytorch/ignite/blob/master/ignite/contrib/metrics/precision_recall_curve.py)
- [Roc Auc](https://github.com/pytorch/ignite/blob/master/ignite/contrib/metrics/roc_auc.py) 
We need to make the implementation simpler. Similar to the one here #1690. New arg `device` will be added, and we need to add the necessary distributed tests as well. See here for the tests https://github.com/pytorch/ignite/blob/master/tests/ignite/contrib/metrics/regression/test_canberra_metric.py#L99



Discussion:

Resolving patch:
diff --git a/ignite/contrib/metrics/average_precision.py b/ignite/contrib/metrics/average_precision.py
--- a/ignite/contrib/metrics/average_precision.py
+++ b/ignite/contrib/metrics/average_precision.py
@@ -1,4 +1,4 @@
-from typing import Callable
+from typing import Callable, Optional, Union
 
 import torch
 
@@ -6,13 +6,10 @@
 
 
 def average_precision_compute_fn(y_preds: torch.Tensor, y_targets: torch.Tensor) -> float:
-    try:
-        from sklearn.metrics import average_precision_score
-    except ImportError:
-        raise RuntimeError("This contrib module requires sklearn to be installed.")
+    from sklearn.metrics import average_precision_score
 
-    y_true = y_targets.numpy()
-    y_pred = y_preds.numpy()
+    y_true = y_targets.cpu().numpy()
+    y_pred = y_preds.cpu().numpy()
     return average_precision_score(y_true, y_pred)
 
 
@@ -30,6 +27,7 @@ class AveragePrecision(EpochMetric):
             <https://scikit-learn.org/stable/modules/generated/sklearn.metrics.average_precision_score.html
             #sklearn.metrics.average_precision_score>`_ is run on the first batch of data to ensure there are
             no issues. User will be warned in case there are any issues computing the function.
+        device: optional device specification for internal storage.
 
     AveragePrecision expects y to be comprised of 0's and 1's. y_pred must either be probability estimates or
     confidence values. To apply an activation to y_pred, use output_transform as shown below:
@@ -45,7 +43,21 @@ def activated_output_transform(output):
 
     """
 
-    def __init__(self, output_transform: Callable = lambda x: x, check_compute_fn: bool = False) -> None:
+    def __init__(
+        self,
+        output_transform: Callable = lambda x: x,
+        check_compute_fn: bool = False,
+        device: Union[str, torch.device] = torch.device("cpu"),
+    ):
+
+        try:
+            from sklearn.metrics import average_precision_score
+        except ImportError:
+            raise RuntimeError("This contrib module requires sklearn to be installed.")
+
         super(AveragePrecision, self).__init__(
-            average_precision_compute_fn, output_transform=output_transform, check_compute_fn=check_compute_fn
+            average_precision_compute_fn,
+            output_transform=output_transform,
+            check_compute_fn=check_compute_fn,
+            device=device,
         )


Test patch:
diff --git a/tests/ignite/contrib/metrics/test_average_precision.py b/tests/ignite/contrib/metrics/test_average_precision.py
--- a/tests/ignite/contrib/metrics/test_average_precision.py
+++ b/tests/ignite/contrib/metrics/test_average_precision.py
@@ -1,111 +1,487 @@
+import os
+
 import numpy as np
+import pytest
 import torch
 from sklearn.metrics import average_precision_score
 
+import ignite.distributed as idist
 from ignite.contrib.metrics import AveragePrecision
 from ignite.engine import Engine
+from ignite.exceptions import NotComputableError
+
+torch.manual_seed(12)
+
+
+def test_no_update():
+    ap = AveragePrecision()
+
+    with pytest.raises(
+        NotComputableError, match=r"EpochMetric must have at least one example before it can be computed"
+    ):
+        ap.compute()
+
+
+def test_input_types():
+    ap = AveragePrecision()
+    ap.reset()
+    output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    ap.update(output1)
+
+    with pytest.raises(ValueError, match=r"Incoherent types between input y_pred and stored predictions"):
+        ap.update((torch.randint(0, 5, size=(4, 3)), torch.randint(0, 2, size=(4, 3))))
+
+    with pytest.raises(ValueError, match=r"Incoherent types between input y and stored targets"):
+        ap.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3)).to(torch.int32)))
+
+    with pytest.raises(ValueError, match=r"Incoherent types between input y_pred and stored predictions"):
+        ap.update((torch.randint(0, 2, size=(10,)).long(), torch.randint(0, 2, size=(10, 5)).long()))
+
+
+def test_check_shape():
+    ap = AveragePrecision()
+
+    with pytest.raises(ValueError, match=r"Predictions should be of shape"):
+        ap._check_shape((torch.tensor(0), torch.tensor(0)))
+
+    with pytest.raises(ValueError, match=r"Predictions should be of shape"):
+        ap._check_shape((torch.rand(4, 3, 1), torch.rand(4, 3)))
+
+    with pytest.raises(ValueError, match=r"Targets should be of shape"):
+        ap._check_shape((torch.rand(4, 3), torch.rand(4, 3, 1)))
+
+
+def test_binary_input_N():
+    ap = AveragePrecision()
+
+    def _test(y_pred, y, n_iters):
+        ap.reset()
+        ap.update((y_pred, y))
+
+        np_y = y.numpy()
+        np_y_pred = y_pred.numpy()
+
+        if n_iters > 1:
+            batch_size = y.shape[0] // n_iters + 1
+            for i in range(n_iters):
+                idx = i * batch_size
+                ap.update((y_pred[idx : idx + batch_size], y[idx : idx + batch_size]))
+
+        res = ap.compute()
+        assert isinstance(res, float)
+        assert average_precision_score(np_y, np_y_pred) == pytest.approx(res)
+
+    test_cases = [
+        (torch.randint(0, 2, size=(10,)).long(), torch.randint(0, 2, size=(10,)).long(), 1),
+        (torch.randint(0, 2, size=(100,)).long(), torch.randint(0, 2, size=(100,)).long(), 1),
+        (torch.randint(0, 2, size=(10, 1)).long(), torch.randint(0, 2, size=(10, 1)).long(), 1),
+        (torch.randint(0, 2, size=(100, 1)).long(), torch.randint(0, 2, size=(100, 1)).long(), 1),
+        # updated batches
+        (torch.randint(0, 2, size=(10,)).long(), torch.randint(0, 2, size=(10,)).long(), 16),
+        (torch.randint(0, 2, size=(100,)).long(), torch.randint(0, 2, size=(100,)).long(), 16),
+        (torch.randint(0, 2, size=(10, 1)).long(), torch.randint(0, 2, size=(10, 1)).long(), 16),
+        (torch.randint(0, 2, size=(100, 1)).long(), torch.randint(0, 2, size=(100, 1)).long(), 16),
+    ]
+
+    for y_pred, y, n_iters in test_cases:
+        # check multiple random inputs as random exact occurencies are rare
+        for _ in range(10):
+            _test(y_pred, y, n_iters)
+
+
+def test_multilabel_input_N():
+
+    ap = AveragePrecision()
+
+    def _test(y_pred, y, n_iters):
+        ap.reset()
+        ap.update((y_pred, y))
+
+        np_y = y.numpy()
+        np_y_pred = y_pred.numpy()
+
+        if n_iters > 1:
+            batch_size = y.shape[0] // n_iters + 1
+            for i in range(n_iters):
+                idx = i * batch_size
+                ap.update((y_pred[idx : idx + batch_size], y[idx : idx + batch_size]))
+
+        res = ap.compute()
+        assert isinstance(res, float)
+        assert average_precision_score(np_y, np_y_pred) == pytest.approx(res)
+
+    test_cases = [
+        (torch.randint(0, 2, size=(10, 4)).long(), torch.randint(0, 2, size=(10, 4)).long(), 1),
+        (torch.randint(0, 2, size=(50, 7)).long(), torch.randint(0, 2, size=(50, 7)).long(), 1),
+        (torch.randint(0, 2, size=(100, 4)).long(), torch.randint(0, 2, size=(100, 4)).long(), 1),
+        (torch.randint(0, 2, size=(200, 6)).long(), torch.randint(0, 2, size=(200, 6)).long(), 1),
+        # updated batches
+        (torch.randint(0, 2, size=(10, 4)).long(), torch.randint(0, 2, size=(10, 4)).long(), 16),
+        (torch.randint(0, 2, size=(50, 7)).long(), torch.randint(0, 2, size=(50, 7)).long(), 16),
+        (torch.randint(0, 2, size=(100, 4)).long(), torch.randint(0, 2, size=(100, 4)).long(), 16),
+        (torch.randint(0, 2, size=(200, 6)).long(), torch.randint(0, 2, size=(200, 6)).long(), 16),
+    ]
+
+    for y_pred, y, n_iters in test_cases:
+        # check multiple random inputs as random exact occurencies are rare
+        for _ in range(10):
+            _test(y_pred, y, n_iters)
+
+
+def test_integration_binary_input_with_output_transform():
+    def _test(y_pred, y, batch_size):
+        def update_fn(engine, batch):
+            idx = (engine.state.iteration - 1) * batch_size
+            y_true_batch = np_y[idx : idx + batch_size]
+            y_pred_batch = np_y_pred[idx : idx + batch_size]
+            return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
+
+        engine = Engine(update_fn)
+
+        ap_metric = AveragePrecision(output_transform=lambda x: (x[1], x[2]))
+        ap_metric.attach(engine, "ap")
+
+        np_y = y.numpy()
+        np_y_pred = y_pred.numpy()
+
+        np_ap = average_precision_score(np_y, np_y_pred)
+
+        data = list(range(y_pred.shape[0] // batch_size))
+        ap = engine.run(data, max_epochs=1).metrics["ap"]
+
+        assert isinstance(ap, float)
+        assert np_ap == pytest.approx(ap)
+
+    test_cases = [
+        (torch.randint(0, 2, size=(100,)).long(), torch.randint(0, 2, size=(100,)).long(), 10),
+        (torch.randint(0, 2, size=(100, 1)).long(), torch.randint(0, 2, size=(100, 1)).long(), 10),
+        (torch.randint(0, 2, size=(200,)).long(), torch.randint(0, 2, size=(200,)).long(), 10),
+        (torch.randint(0, 2, size=(200, 1)).long(), torch.randint(0, 2, size=(200, 1)).long(), 10),
+    ]
+
+    for y_pred, y, batch_size in test_cases:
+        # check multiple random inputs as random exact occurencies are rare
+        for _ in range(10):
+            _test(y_pred, y, batch_size)
+
+
+def test_integration_multilabel_input_with_output_transform():
+    def _test(y_pred, y, batch_size):
+        def update_fn(engine, batch):
+            idx = (engine.state.iteration - 1) * batch_size
+            y_true_batch = np_y[idx : idx + batch_size]
+            y_pred_batch = np_y_pred[idx : idx + batch_size]
+            return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
+
+        engine = Engine(update_fn)
+
+        ap_metric = AveragePrecision(output_transform=lambda x: (x[1], x[2]))
+        ap_metric.attach(engine, "ap")
+
+        np_y = y.numpy()
+        np_y_pred = y_pred.numpy()
+
+        np_ap = average_precision_score(np_y, np_y_pred)
+
+        data = list(range(y_pred.shape[0] // batch_size))
+        ap = engine.run(data, max_epochs=1).metrics["ap"]
+
+        assert isinstance(ap, float)
+        assert np_ap == pytest.approx(ap)
+
+    test_cases = [
+        (torch.randint(0, 2, size=(100, 3)).long(), torch.randint(0, 2, size=(100, 3)).long(), 10),
+        (torch.randint(0, 2, size=(100, 4)).long(), torch.randint(0, 2, size=(100, 4)).long(), 10),
+        (torch.randint(0, 2, size=(200, 5)).long(), torch.randint(0, 2, size=(200, 5)).long(), 10),
+        (torch.randint(0, 2, size=(200, 6)).long(), torch.randint(0, 2, size=(200, 6)).long(), 10),
+    ]
+
+    for y_pred, y, batch_size in test_cases:
+        # check multiple random inputs as random exact occurencies are rare
+        for _ in range(10):
+            _test(y_pred, y, batch_size)
+
+
+def _test_distirb_binary_input_N(device):
+
+    rank = idist.get_rank()
+    torch.manual_seed(12)
+
+    def _test(y_pred, y, n_iters, metric_device):
+
+        metric_device = torch.device(metric_device)
+        ap = AveragePrecision(device=metric_device)
+        torch.manual_seed(10 + rank)
+
+        ap.reset()
+        ap.update((y_pred, y))
+
+        if n_iters > 1:
+            batch_size = y.shape[0] // n_iters + 1
+            for i in range(n_iters):
+                idx = i * batch_size
+                ap.update((y_pred[idx : idx + batch_size], y[idx : idx + batch_size]))
+
+        # gather y_pred, y
+        y_pred = idist.all_gather(y_pred)
+        y = idist.all_gather(y)
+
+        np_y = y.cpu().numpy()
+        np_y_pred = y_pred.cpu().numpy()
+
+        res = ap.compute()
+        assert isinstance(res, float)
+        assert average_precision_score(np_y, np_y_pred) == pytest.approx(res)
+
+    test_cases = [
+        (torch.randint(0, 2, size=(10,)).long(), torch.randint(0, 2, size=(10,)).long(), 1),
+        (torch.randint(0, 2, size=(100,)).long(), torch.randint(0, 2, size=(100,)).long(), 1),
+        (torch.randint(0, 2, size=(10, 1)).long(), torch.randint(0, 2, size=(10, 1)).long(), 1),
+        (torch.randint(0, 2, size=(100, 1)).long(), torch.randint(0, 2, size=(100, 1)).long(), 1),
+        # updated batches
+        (torch.randint(0, 2, size=(10,)).long(), torch.randint(0, 2, size=(10,)).long(), 16),
+        (torch.randint(0, 2, size=(100,)).long(), torch.randint(0, 2, size=(100,)).long(), 16),
+        (torch.randint(0, 2, size=(10, 1)).long(), torch.randint(0, 2, size=(10, 1)).long(), 16),
+        (torch.randint(0, 2, size=(100, 1)).long(), torch.randint(0, 2, size=(100, 1)).long(), 16),
+    ]
+
+    for y_pred, y, batch_size in test_cases:
+        for _ in range(3):
+            _test(y_pred, y, batch_size, "cpu")
+            if device.type != "xla":
+                _test(y_pred, y, batch_size, idist.device())
+
+
+def _test_distirb_multilabel_input_N(device):
+
+    rank = idist.get_rank()
+    torch.manual_seed(12)
+
+    def _test(y_pred, y, n_iters, metric_device):
+
+        metric_device = torch.device(metric_device)
+        ap = AveragePrecision(device=metric_device)
+        torch.manual_seed(10 + rank)
+
+        ap.reset()
+        ap.update((y_pred, y))
+
+        if n_iters > 1:
+            batch_size = y.shape[0] // n_iters + 1
+            for i in range(n_iters):
+                idx = i * batch_size
+                ap.update((y_pred[idx : idx + batch_size], y[idx : idx + batch_size]))
+
+        # gather y_pred, y
+        y_pred = idist.all_gather(y_pred)
+        y = idist.all_gather(y)
+
+        np_y = y.cpu().numpy()
+        np_y_pred = y_pred.cpu().numpy()
+
+        res = ap.compute()
+        assert isinstance(res, float)
+        assert average_precision_score(np_y, np_y_pred) == pytest.approx(res)
+
+    test_cases = [
+        (torch.randint(0, 2, size=(10, 4)).long(), torch.randint(0, 2, size=(10, 4)).long(), 1),
+        (torch.randint(0, 2, size=(100, 7)).long(), torch.randint(0, 2, size=(100, 7)).long(), 1),
+        (torch.randint(0, 2, size=(100, 5)).long(), torch.randint(0, 2, size=(100, 5)).long(), 1),
+        (torch.randint(0, 2, size=(100, 3)).long(), torch.randint(0, 2, size=(100, 3)).long(), 1),
+        # updated batches
+        (torch.randint(0, 2, size=(10, 4)).long(), torch.randint(0, 2, size=(10, 4)).long(), 16),
+        (torch.randint(0, 2, size=(100, 7)).long(), torch.randint(0, 2, size=(100, 7)).long(), 16),
+        (torch.randint(0, 2, size=(100, 5)).long(), torch.randint(0, 2, size=(100, 5)).long(), 16),
+        (torch.randint(0, 2, size=(100, 3)).long(), torch.randint(0, 2, size=(100, 3)).long(), 16),
+    ]
+
+    for y_pred, y, batch_size in test_cases:
+        for _ in range(3):
+            _test(y_pred, y, batch_size, "cpu")
+            if device.type != "xla":
+                _test(y_pred, y, batch_size, idist.device())
+
+
+def _test_distrib_integration_binary(device):
+
+    rank = idist.get_rank()
+    torch.manual_seed(12)
+
+    def _test(n_epochs, metric_device):
+        metric_device = torch.device(metric_device)
+        n_iters = 80
+        s = 16
+        n_classes = 2
+
+        offset = n_iters * s
+        y_true = torch.randint(0, n_classes, size=(offset * idist.get_world_size(),)).to(device)
+        y_preds = torch.rand(offset * idist.get_world_size(),).to(device)
+
+        def update(engine, i):
+            return (
+                y_preds[i * s + rank * offset : (i + 1) * s + rank * offset],
+                y_true[i * s + rank * offset : (i + 1) * s + rank * offset],
+            )
+
+        engine = Engine(update)
+
+        ap = AveragePrecision(device=metric_device)
+        ap.attach(engine, "ap")
+
+        data = list(range(n_iters))
+        engine.run(data=data, max_epochs=n_epochs)
+
+        assert "ap" in engine.state.metrics
+
+        res = engine.state.metrics["ap"]
+        if isinstance(res, torch.Tensor):
+            res = res.cpu().numpy()
+
+        true_res = average_precision_score(y_true.cpu().numpy(), y_preds.cpu().numpy())
+
+        assert pytest.approx(res) == true_res
+
+    metric_devices = ["cpu"]
+    if device.type != "xla":
+        metric_devices.append(idist.device())
+    for metric_device in metric_devices:
+        for _ in range(2):
+            _test(n_epochs=1, metric_device=metric_device)
+            _test(n_epochs=2, metric_device=metric_device)
+
+
+def _test_distrib_integration_multilabel(device):
+
+    rank = idist.get_rank()
+    torch.manual_seed(12)
+
+    def _test(n_epochs, metric_device):
+        metric_device = torch.device(metric_device)
+        n_iters = 80
+        s = 16
+        n_classes = 2
 
+        offset = n_iters * s
+        y_true = torch.randint(0, n_classes, size=(offset * idist.get_world_size(), 10)).to(device)
+        y_preds = torch.randint(0, n_classes, size=(offset * idist.get_world_size(), 10)).to(device)
 
-def test_ap_score():
+        def update(engine, i):
+            return (
+                y_preds[i * s + rank * offset : (i + 1) * s + rank * offset, :],
+                y_true[i * s + rank * offset : (i + 1) * s + rank * offset, :],
+            )
 
-    size = 100
-    np_y_pred = np.random.rand(size, 5)
-    np_y = np.random.randint(0, 2, size=(size, 5), dtype=np.long)
-    np_ap = average_precision_score(np_y, np_y_pred)
+        engine = Engine(update)
 
-    ap_metric = AveragePrecision()
-    y_pred = torch.from_numpy(np_y_pred)
-    y = torch.from_numpy(np_y)
+        ap = AveragePrecision(device=metric_device)
+        ap.attach(engine, "ap")
 
-    ap_metric.reset()
-    ap_metric.update((y_pred, y))
-    ap = ap_metric.compute()
+        data = list(range(n_iters))
+        engine.run(data=data, max_epochs=n_epochs)
 
-    assert ap == np_ap
+        assert "ap" in engine.state.metrics
 
+        res = engine.state.metrics["ap"]
+        if isinstance(res, torch.Tensor):
+            res = res.cpu().numpy()
 
-def test_ap_score_2():
+        true_res = average_precision_score(y_true.cpu().numpy(), y_preds.cpu().numpy())
 
-    np.random.seed(1)
-    size = 100
-    np_y_pred = np.random.rand(size, 1)
-    np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2 :] = 1
-    np.random.shuffle(np_y)
-    np_ap = average_precision_score(np_y, np_y_pred)
+        assert pytest.approx(res) == true_res
 
-    ap_metric = AveragePrecision()
-    y_pred = torch.from_numpy(np_y_pred)
-    y = torch.from_numpy(np_y)
+    metric_devices = ["cpu"]
+    if device.type != "xla":
+        metric_devices.append(idist.device())
+    for metric_device in metric_devices:
+        for _ in range(2):
+            _test(n_epochs=1, metric_device=metric_device)
+            _test(n_epochs=2, metric_device=metric_device)
 
-    ap_metric.reset()
-    n_iters = 10
-    batch_size = size // n_iters
-    for i in range(n_iters):
-        idx = i * batch_size
-        ap_metric.update((y_pred[idx : idx + batch_size], y[idx : idx + batch_size]))
 
-    ap = ap_metric.compute()
+@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(distributed_context_single_node_nccl):
 
-    assert ap == np_ap
+    device = torch.device(f"cuda:{distributed_context_single_node_nccl['local_rank']}")
+    _test_distirb_binary_input_N(device)
+    _test_distirb_multilabel_input_N(device)
+    _test_distrib_integration_binary(device)
+    _test_distrib_integration_multilabel(device)
 
 
-def test_integration_ap_score_with_output_transform():
+@pytest.mark.distributed
+@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):
 
-    np.random.seed(1)
-    size = 100
-    np_y_pred = np.random.rand(size, 1)
-    np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2 :] = 1
-    np.random.shuffle(np_y)
+    device = torch.device("cpu")
+    _test_distirb_binary_input_N(device)
+    _test_distirb_multilabel_input_N(device)
+    _test_distrib_integration_binary(device)
+    _test_distrib_integration_multilabel(device)
 
-    np_ap = average_precision_score(np_y, np_y_pred)
 
-    batch_size = 10
+@pytest.mark.distributed
+@pytest.mark.skipif(not idist.has_hvd_support, reason="Skip if no Horovod dist support")
+@pytest.mark.skipif("WORLD_SIZE" in os.environ, reason="Skip if launched as multiproc")
+def test_distrib_hvd(gloo_hvd_executor):
 
-    def update_fn(engine, batch):
-        idx = (engine.state.iteration - 1) * batch_size
-        y_true_batch = np_y[idx : idx + batch_size]
-        y_pred_batch = np_y_pred[idx : idx + batch_size]
-        return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
+    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()
 
-    engine = Engine(update_fn)
+    gloo_hvd_executor(_test_distirb_binary_input_N, (device,), np=nproc, do_init=True)
+    gloo_hvd_executor(_test_distirb_multilabel_input_N, (device,), np=nproc, do_init=True)
+    gloo_hvd_executor(_test_distrib_integration_binary, (device,), np=nproc, do_init=True)
+    gloo_hvd_executor(_test_distrib_integration_multilabel, (device,), np=nproc, do_init=True)
 
-    ap_metric = AveragePrecision(output_transform=lambda x: (x[1], x[2]))
-    ap_metric.attach(engine, "ap")
 
-    data = list(range(size // batch_size))
-    ap = engine.run(data, max_epochs=1).metrics["ap"]
+@pytest.mark.multinode_distributed
+@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):
 
-    assert ap == np_ap
+    device = torch.device("cpu")
+    _test_distirb_binary_input_N(device)
+    _test_distirb_multilabel_input_N(device)
+    _test_distrib_integration_binary(device)
+    _test_distrib_integration_multilabel(device)
 
 
-def test_integration_ap_score_with_activated_output_transform():
+@pytest.mark.multinode_distributed
+@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):
 
-    np.random.seed(1)
-    size = 100
-    np_y_pred = np.random.rand(size, 1)
-    np_y_pred_softmax = torch.softmax(torch.from_numpy(np_y_pred), dim=1).numpy()
-    np_y = np.zeros((size,), dtype=np.long)
-    np_y[size // 2 :] = 1
-    np.random.shuffle(np_y)
+    device = torch.device(f"cuda:{distributed_context_multi_node_nccl['local_rank']}")
+    _test_distirb_binary_input_N(device)
+    _test_distirb_multilabel_input_N(device)
+    _test_distrib_integration_binary(device)
+    _test_distrib_integration_multilabel(device)
 
-    np_ap = average_precision_score(np_y, np_y_pred_softmax)
 
-    batch_size = 10
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" in os.environ, reason="Skip if NUM_TPU_WORKERS is in env vars")
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_distrib_single_device_xla():
 
-    def update_fn(engine, batch):
-        idx = (engine.state.iteration - 1) * batch_size
-        y_true_batch = np_y[idx : idx + batch_size]
-        y_pred_batch = np_y_pred[idx : idx + batch_size]
-        return idx, torch.from_numpy(y_pred_batch), torch.from_numpy(y_true_batch)
+    device = idist.device()
+    _test_distirb_binary_input_N(device)
+    _test_distirb_multilabel_input_N(device)
+    _test_distrib_integration_binary(device)
+    _test_distrib_integration_multilabel(device)
 
-    engine = Engine(update_fn)
 
-    ap_metric = AveragePrecision(output_transform=lambda x: (torch.softmax(x[1], dim=1), x[2]))
-    ap_metric.attach(engine, "ap")
+def _test_distrib_xla_nprocs(index):
+    device = idist.device()
+    _test_distirb_binary_input_N(device)
+    _test_distirb_multilabel_input_N(device)
+    _test_distrib_integration_binary(device)
+    _test_distrib_integration_multilabel(device)
 
-    data = list(range(size // batch_size))
-    ap = engine.run(data, max_epochs=1).metrics["ap"]
 
-    assert ap == np_ap
+@pytest.mark.tpu
+@pytest.mark.skipif("NUM_TPU_WORKERS" not in os.environ, reason="Skip if no NUM_TPU_WORKERS in env vars")
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_distrib_xla_nprocs(xmp_executor):
+    n = int(os.environ["NUM_TPU_WORKERS"])
+    xmp_executor(_test_distrib_xla_nprocs, args=(), nprocs=n)