Support flexible EpochMetric compute_fn outputs

ignite/metrics/epoch_metric.py

@@ -1,6 +1,6 @@
 import warnings
-from collections.abc import Callable
-from typing import cast
+from collections.abc import Callable, Mapping, Sequence
+from typing import cast, Union
 
 import torch
 
@@ -10,6 +10,9 @@
 
 __all__ = ["EpochMetric"]
 
+# Supported return types for ``EpochMetric``'s ``compute_fn``.
+EpochMetricOutput = Union[int, float, torch.Tensor, Sequence, Mapping]
+
 
 class EpochMetric(Metric):
     """Class for metrics that should be computed on the entire output history of a model.
@@ -30,7 +33,13 @@ class EpochMetric(Metric):
 
     Args:
         compute_fn: a callable which receives two tensors as the `predictions` and `targets`
-            and returns a scalar. Input tensors will be on specified ``device`` (see arg below).
+            and returns the computed metric. Supported return types are: ``int``, ``float``,
+            ``torch.Tensor``, a ``Sequence`` (tuple/list) of these, or a ``Mapping`` (dict) with
+            string keys and these values. An unsupported return type raises a ``TypeError``.
+            Note: in distributed configuration (``world_size > 1``), only scalar and
+            ``torch.Tensor`` outputs are broadcast across processes; tuple/list/mapping outputs
+            are supported only when ``world_size == 1``. Input tensors will be on specified
+            ``device`` (see arg below).
         output_transform: a callable that is used to transform the
             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
@@ -93,7 +102,7 @@ def __init__(
     def reset(self) -> None:
         self._predictions: list[torch.Tensor] = []
         self._targets: list[torch.Tensor] = []
-        self._result: float | None = None
+        self._result: EpochMetricOutput | None = None
 
     def _check_shape(self, output: tuple[torch.Tensor, torch.Tensor]) -> None:
         y_pred, y = output
@@ -142,7 +151,27 @@ def update(self, output: tuple[torch.Tensor, torch.Tensor]) -> None:
             except Exception as e:
                 warnings.warn(f"Probably, there can be a problem with `compute_fn`:\n {e}.", EpochMetricWarning)
 
-    def compute(self) -> float:
+    def _check_output_type(self, result: EpochMetricOutput) -> None:
+        # Recursively validate that compute_fn's output is a supported type. ``str``/``bytes``
+        # are rejected explicitly since ``str`` is itself a ``Sequence``.
+        if isinstance(result, (int, float, torch.Tensor)):
+            return
+        if isinstance(result, Mapping):
+            for key, value in result.items():
+                if not isinstance(key, str):
+                    raise TypeError(f"compute_fn output mapping keys should be str, but given {type(key)}.")
+                self._check_output_type(value)
+            return
+        if isinstance(result, Sequence) and not isinstance(result, (str, bytes)):
+            for value in result:
+                self._check_output_type(value)
+            return
+        raise TypeError(
+            f"compute_fn output type {type(result)} is not supported. Supported types are: "
+            "int, float, torch.Tensor, a Sequence of these, or a Mapping with str keys and these values."
+        )
+
+    def compute(self) -> EpochMetricOutput:
         if len(self._predictions) < 1 or len(self._targets) < 1:
             raise NotComputableError(f"{type(self).__name__} must have at least one example before it can be computed.")
 
@@ -156,14 +185,48 @@ def compute(self) -> float:
                 _prediction_tensor = cast(torch.Tensor, idist.all_gather(_prediction_tensor))
                 _target_tensor = cast(torch.Tensor, idist.all_gather(_target_tensor))
 
-            self._result = 0.0
+            result: EpochMetricOutput = 0.0
             if idist.get_rank() == 0:
                 # Run compute_fn on zero rank only
-                self._result = self.compute_fn(_prediction_tensor, _target_tensor)
+                result = self.compute_fn(_prediction_tensor, _target_tensor)
 
             if ws > 1:
-                # broadcast result to all processes
-                self._result = cast(float, idist.broadcast(self._result, src=0))
+                # All ranks must take the same path through the collective calls below, otherwise
+                # they would deadlock on mismatched broadcasts. Only rank 0 holds the real result,
+                # so it classifies the result and shares a status code with every rank *before*
+                # broadcasting the result itself. Type/validation problems are surfaced through the
+                # same mechanism so that every rank raises the same exception.
+                _BROADCASTABLE, _UNSUPPORTED_CONTAINER, _UNSUPPORTED_TYPE = 0, 1, 2
+                status = _BROADCASTABLE
+                if idist.get_rank() == 0 and not isinstance(result, (int, float, torch.Tensor)):
+                    try:
+                        self._check_output_type(result)
+                        status = _UNSUPPORTED_CONTAINER
+                    except TypeError:
+                        status = _UNSUPPORTED_TYPE
+                status = int(idist.broadcast(status, src=0))
+
+                if status == _UNSUPPORTED_TYPE:
+                    # Every rank raises the same error; no result broadcast is attempted.
+                    raise TypeError(
+                        "compute_fn output type is not supported. Supported types are: int, float, "
+                        "torch.Tensor, a Sequence of these, or a Mapping with str keys and these values."
+                    )
+                if status == _UNSUPPORTED_CONTAINER:
+                    # Every rank raises the same error; no result broadcast is attempted.
+                    raise NotImplementedError(
+                        "Distributed broadcast of tuple/list/mapping compute_fn outputs is not supported yet. "
+                        "Such outputs are currently supported only in non-distributed (world_size == 1) "
+                        "configuration."
+                    )
+
+                # status == _BROADCASTABLE: every rank performs the matching result broadcast.
+                result = cast(EpochMetricOutput, idist.broadcast(result, src=0, safe_mode=True))
+            else:
+                # Single process: validate directly and surface unsupported types as TypeError.
+                self._check_output_type(result)
+
+            self._result = result
 
         return self._result
 

tests/ignite/metrics/test_epoch_metric.py

@@ -211,3 +211,234 @@ def compute_fn(y_preds, y_targets):
     assert torch.equal(em._targets[0].cpu(), output1[1].cpu())
     assert torch.equal(em._targets[1].cpu(), output2[1].cpu())
     assert em.compute() == 0.0
+
+
+def test_epoch_metric_scalar_tensor_output(available_device):
+    # compute_fn returns a 0-dim (scalar) tensor instead of a python float
+    def compute_fn(y_preds, y_targets):
+        return torch.mean(y_preds - y_targets.type_as(y_preds))
+
+    em = EpochMetric(compute_fn, device=available_device)
+
+    em.reset()
+    output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output1)
+    output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output2)
+
+    preds = torch.cat([output1[0], output2[0]], dim=0)
+    targets = torch.cat([output1[1], output2[1]], dim=0)
+    expected = compute_fn(preds, targets)
+
+    result = em.compute()
+    assert isinstance(result, torch.Tensor)
+    assert result.ndim == 0
+    assert torch.allclose(result.cpu(), expected.cpu())
+
+
+def test_epoch_metric_vector_tensor_output(available_device):
+    # compute_fn returns a 1-dim (vector) tensor, one value per target column
+    def compute_fn(y_preds, y_targets):
+        return y_preds.mean(dim=0)
+
+    em = EpochMetric(compute_fn, device=available_device)
+
+    em.reset()
+    output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output1)
+    output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output2)
+
+    preds = torch.cat([output1[0], output2[0]], dim=0)
+    targets = torch.cat([output1[1], output2[1]], dim=0)
+    expected = compute_fn(preds, targets)
+
+    result = em.compute()
+    assert isinstance(result, torch.Tensor)
+    assert result.shape == (3,)
+    assert torch.allclose(result.cpu(), expected.cpu())
+
+
+def test_epoch_metric_tuple_of_tensors_output(available_device):
+    # compute_fn returns a tuple of tensors
+    def compute_fn(y_preds, y_targets):
+        return y_preds.mean(dim=0), y_preds.sum(dim=0)
+
+    em = EpochMetric(compute_fn, device=available_device)
+
+    em.reset()
+    output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output1)
+    output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output2)
+
+    preds = torch.cat([output1[0], output2[0]], dim=0)
+    targets = torch.cat([output1[1], output2[1]], dim=0)
+    expected = compute_fn(preds, targets)
+
+    result = em.compute()
+    assert isinstance(result, tuple)
+    assert len(result) == 2
+    for r, e in zip(result, expected):
+        assert torch.allclose(r.cpu(), e.cpu())
+
+
+def test_epoch_metric_list_of_tensors_output(available_device):
+    # compute_fn returns a list of tensors
+    def compute_fn(y_preds, y_targets):
+        return [y_preds.mean(dim=0), y_preds.sum(dim=0)]
+
+    em = EpochMetric(compute_fn, device=available_device)
+
+    em.reset()
+    output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output1)
+    output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output2)
+
+    preds = torch.cat([output1[0], output2[0]], dim=0)
+    targets = torch.cat([output1[1], output2[1]], dim=0)
+    expected = compute_fn(preds, targets)
+
+    result = em.compute()
+    assert isinstance(result, list)
+    assert len(result) == 2
+    for r, e in zip(result, expected):
+        assert torch.allclose(r.cpu(), e.cpu())
+
+
+def test_epoch_metric_mapping_of_tensors_output(available_device):
+    # compute_fn returns a dict mapping str -> tensor
+    def compute_fn(y_preds, y_targets):
+        return {"mean": y_preds.mean(dim=0), "sum": y_preds.sum(dim=0)}
+
+    em = EpochMetric(compute_fn, device=available_device)
+
+    em.reset()
+    output1 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output1)
+    output2 = (torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long))
+    em.update(output2)
+
+    preds = torch.cat([output1[0], output2[0]], dim=0)
+    targets = torch.cat([output1[1], output2[1]], dim=0)
+    expected = compute_fn(preds, targets)
+
+    result = em.compute()
+    assert isinstance(result, dict)
+    assert set(result.keys()) == {"mean", "sum"}
+    for key in expected:
+        assert torch.allclose(result[key].cpu(), expected[key].cpu())
+
+
+def test_epoch_metric_unsupported_output_type_raises(available_device):
+    # An unsupported output type (str) should raise a clear TypeError.
+    def compute_fn(y_preds, y_targets):
+        return "not-a-number"
+
+    em = EpochMetric(compute_fn, check_compute_fn=False, device=available_device)
+
+    em.reset()
+    em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)))
+    em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)))
+
+    with pytest.raises(TypeError, match=r"compute_fn output type .* is not supported"):
+        em.compute()
+
+
+def test_epoch_metric_nested_invalid_output_raises(available_device):
+    # A container holding an unsupported value (str) should be rejected by the recursive check.
+    def compute_fn(y_preds, y_targets):
+        return [torch.tensor(1.0), "not-a-number"]
+
+    em = EpochMetric(compute_fn, check_compute_fn=False, device=available_device)
+
+    em.reset()
+    em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)))
+    em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)))
+
+    with pytest.raises(TypeError, match=r"compute_fn output type .* is not supported"):
+        em.compute()
+
+
+def test_epoch_metric_mapping_non_str_key_raises(available_device):
+    # A mapping with a non-string key is not a supported output.
+    def compute_fn(y_preds, y_targets):
+        return {0: torch.tensor(1.0)}
+
+    em = EpochMetric(compute_fn, check_compute_fn=False, device=available_device)
+
+    em.reset()
+    em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)))
+    em.update((torch.rand(4, 3), torch.randint(0, 2, size=(4, 3), dtype=torch.long)))
+
+    with pytest.raises(TypeError, match=r"mapping keys should be str"):
+        em.compute()
+
+
+def test_distrib_output_tensor(distributed):
+    # A non-scalar tensor output should broadcast successfully and be identical on all ranks.
+    device = idist.device() if idist.device().type != "xla" else "cpu"
+    torch.manual_seed(40 + idist.get_rank())
+
+    def compute_fn(y_preds, y_targets):
+        return y_preds.mean(dim=0)
+
+    em = EpochMetric(compute_fn, check_compute_fn=False, device=device)
+    em.reset()
+    em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device)))
+    em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device)))
+
+    result = em.compute()
+    assert isinstance(result, torch.Tensor)
+    assert result.shape == (3,)
+    # All ranks must receive the same broadcasted values.
+    gathered = idist.all_gather(result.unsqueeze(0))
+    assert torch.allclose(gathered[0], gathered[-1])
+
+
+@pytest.mark.parametrize(
+    "compute_fn",
+    [
+        lambda y_preds, y_targets: (y_preds.mean(dim=0), y_preds.sum(dim=0)),
+        lambda y_preds, y_targets: [y_preds.mean(dim=0), y_preds.sum(dim=0)],
+        lambda y_preds, y_targets: {"mean": y_preds.mean(dim=0), "sum": y_preds.sum(dim=0)},
+    ],
+    ids=["tuple", "list", "dict"],
+)
+def test_distrib_output_container_raises(distributed, compute_fn):
+    # Documents the current conservative distributed behavior: tuple/list/mapping outputs are
+    # supported in single-process mode, but under world_size > 1 they intentionally fail
+    # symmetrically with NotImplementedError rather than entering mismatched collective calls.
+    # This is a safety/regression guarantee for this implementation, not the ideal long-term
+    # feature behavior (broadcasting containers across ranks is left for future work).
+    if idist.get_world_size() == 1:
+        pytest.skip("This test verifies world_size > 1 behavior; containers are supported in single process.")
+
+    device = idist.device() if idist.device().type != "xla" else "cpu"
+    torch.manual_seed(40 + idist.get_rank())
+
+    em = EpochMetric(compute_fn, check_compute_fn=False, device=device)
+    em.reset()
+    em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device)))
+    em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device)))
+
+    with pytest.raises(NotImplementedError, match=r"Distributed broadcast of tuple/list/mapping"):
+        em.compute()
+
+
+def test_distrib_output_unsupported_raises(distributed):
+    # An unsupported output type (str) should raise TypeError consistently on all ranks.
+    device = idist.device() if idist.device().type != "xla" else "cpu"
+    torch.manual_seed(40 + idist.get_rank())
+
+    def compute_fn(y_preds, y_targets):
+        return "not-a-number"
+
+    em = EpochMetric(compute_fn, check_compute_fn=False, device=device)
+    em.reset()
+    em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device)))
+    em.update((torch.rand(4, 3, device=device), torch.randint(0, 2, size=(4, 3), dtype=torch.long, device=device)))
+
+    with pytest.raises(TypeError, match=r"compute_fn output type.*is not supported"):
+        em.compute()