feat: extend EpochMetric to support more compute_fn output types

ignite/metrics/clustering/_base.py

@@ -1,3 +1,5 @@
+from typing import cast
+
 from torch import Tensor
 from ignite.exceptions import NotComputableError
 from ignite.metrics.epoch_metric import EpochMetric
@@ -36,4 +38,4 @@ def compute(self) -> float:
                 f"{self.__class__.__name__} must have at least one example before it can be computed."
             )
 
-        return super().compute()
+        return cast(float, super().compute())

ignite/metrics/clustering/calinski_harabasz_score.py

@@ -20,7 +20,7 @@ def _calinski_harabasz_score(features: Tensor, labels: Tensor) -> float:
     np_features = features.cpu().numpy()
     np_labels = labels.cpu().numpy()
     score = calinski_harabasz_score(np_features, np_labels)
-    return score
+    return float(score)
 
 
 class CalinskiHarabaszScore(_ClusteringMetricBase):

ignite/metrics/clustering/davies_bouldin_score.py

@@ -20,7 +20,7 @@ def _davies_bouldin_score(features: Tensor, labels: Tensor) -> float:
     np_features = features.cpu().numpy()
     np_labels = labels.cpu().numpy()
     score = davies_bouldin_score(np_features, np_labels)
-    return score
+    return float(score)
 
 
 class DaviesBouldinScore(_ClusteringMetricBase):

ignite/metrics/clustering/silhouette_score.py

@@ -127,4 +127,4 @@ def _silhouette_score(self, features: Tensor, labels: Tensor) -> float:
         np_labels = labels.detach().cpu().numpy()
 
         score = silhouette_score(np_features, np_labels, **self._silhouette_kwargs)
-        return score
+        return float(score)

ignite/metrics/epoch_metric.py

@@ -1,15 +1,29 @@
 import warnings
-from collections.abc import Callable
-from typing import cast
+from collections.abc import Callable, Mapping, Sequence
+from typing import Union, cast
 
 import torch
 
 import ignite.distributed as idist
 from ignite.exceptions import NotComputableError
 from ignite.metrics.metric import Metric, reinit__is_reduced
 
+# Supported return types for ``EpochMetric``'s ``compute_fn``.
+EpochMetricOutput = Union[int, float, torch.Tensor, Sequence, Mapping]
+
 __all__ = ["EpochMetric"]
 
+# Type tags used by ``EpochMetric._broadcast_result`` to let every rank agree on how to
+# decode the next value(s) coming out of the collective calls. ``_TAG_UNSUPPORTED`` is also
+# reused to reject non-str mapping keys, since both cases mean "every rank should raise".
+_TAG_INT = 0
+_TAG_FLOAT = 1
+_TAG_TENSOR = 2
+_TAG_TUPLE = 3
+_TAG_LIST = 4
+_TAG_MAPPING = 5
+_TAG_UNSUPPORTED = -1
+
 
 class EpochMetric(Metric):
     """Class for metrics that should be computed on the entire output history of a model.
@@ -30,7 +44,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, including arbitrarily nested combinations of these.
+            An unsupported return type raises a ``TypeError``. These types are also supported
+            in distributed configuration (``world_size > 1``): the result is broadcast from
+            rank 0 to all other ranks. 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 +113,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 +162,114 @@ 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 _broadcast_result(self, result: EpochMetricOutput, src: int = 0) -> EpochMetricOutput:
+        """Recursively broadcast compute_fn output from src rank to all ranks.
+
+        Each step only broadcasts types that ``idist.broadcast`` natively supports
+        (int, float, torch.Tensor, str), so containers are transmitted by first
+        synchronising their structure (type tag, length, dict keys) and then
+        broadcasting each leaf individually.
+
+        Every rank must take the same path through the collective calls below, so any
+        rejection (unsupported type, non-str mapping key) is decided from a value that has
+        already been broadcast to all ranks, never from a check that only src has run. This
+        way every rank raises together instead of some ranks hanging on a broadcast that src
+        never issues.
+        """
+        rank = idist.get_rank()
+
+        # Step 1: broadcast type tag so all ranks know what to expect
+        if rank == src:
+            if isinstance(result, int):
+                tag = _TAG_INT
+            elif isinstance(result, float):
+                tag = _TAG_FLOAT
+            elif isinstance(result, torch.Tensor):
+                tag = _TAG_TENSOR
+            elif isinstance(result, tuple):
+                tag = _TAG_TUPLE
+            elif isinstance(result, list):
+                tag = _TAG_LIST
+            elif isinstance(result, Mapping):
+                tag = _TAG_MAPPING
+            else:
+                tag = _TAG_UNSUPPORTED
+        else:
+            tag = _TAG_INT
+        tag = cast(int, idist.broadcast(tag, src=src))
+
+        if tag == _TAG_UNSUPPORTED:
+            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."
+            )
+
+        # Step 2: broadcast content based on type. ``result`` is only meaningfully typed on
+        # ``src`` (the tag protocol above guarantees every rank agrees on which branch runs),
+        # so it is cast to the type the tag promises; non-src ranks only need a same-typed
+        # placeholder since their value is discarded by the collective call.
+        if tag == _TAG_INT:
+            int_value = cast(int, result) if rank == src else 0
+            return cast(int, idist.broadcast(int_value, src=src))
+        if tag == _TAG_FLOAT:
+            float_value = cast(float, result) if rank == src else 0.0
+            return cast(float, idist.broadcast(float_value, src=src))
+        if tag == _TAG_TENSOR:
+            tensor_value = cast(torch.Tensor, result) if rank == src else None
+            return cast(torch.Tensor, idist.broadcast(tensor_value, src=src, safe_mode=True))
+
+        if tag in (_TAG_TUPLE, _TAG_LIST):
+            seq_value = cast(Sequence, result) if rank == src else []
+            length = cast(int, idist.broadcast(len(seq_value) if rank == src else 0, src=src))
+            elements = []
+            for i in range(length):
+                elem = cast(EpochMetricOutput, seq_value[i]) if rank == src else 0
+                elements.append(self._broadcast_result(elem, src=src))
+            return tuple(elements) if tag == _TAG_TUPLE else elements
+
+        # tag == _TAG_MAPPING
+        mapping_value = cast(Mapping, result) if rank == src else {}
+        src_keys = list(mapping_value.keys()) if rank == src else []
+        n_keys = cast(int, idist.broadcast(len(src_keys) if rank == src else 0, src=src))
+        keys = []
+        for i in range(n_keys):
+            raw_key = src_keys[i] if rank == src else None
+            # Broadcast whether this key is a valid (str) key before broadcasting the key
+            # itself: src and non-src ranks must agree on the wire type (str) they are about
+            # to exchange, so this can't be decided from a src-only isinstance check.
+            key_valid = cast(int, idist.broadcast(0 if (rank != src or isinstance(raw_key, str)) else -1, src=src))
+            if key_valid == _TAG_UNSUPPORTED:
+                detail = f" but given {type(raw_key)}" if rank == src else ""
+                raise TypeError(f"compute_fn output mapping keys should be str{detail}.")
+            key = cast(str, idist.broadcast(raw_key if rank == src else "", src=src))
+            keys.append(key)
+        values = []
+        for i in range(n_keys):
+            val = cast(EpochMetricOutput, mapping_value[keys[i]]) if rank == src else 0
+            values.append(self._broadcast_result(val, src=src))
+        return dict(zip(keys, 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 +283,22 @@ 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))
+                # Type/key validation happens inside `_broadcast_result` itself (via the tag
+                # protocol), so every rank reaches the same TypeError together. Do not
+                # pre-validate on rank 0 alone here: that would let rank 0 raise and return
+                # before issuing the first broadcast, leaving other ranks waiting on a
+                # collective call rank 0 never makes.
+                result = self._broadcast_result(result, src=0)
+            else:
+                self._check_output_type(result)
+
+            self._result = result
 
         return self._result
 

ignite/metrics/regression/kendall_correlation.py

@@ -1,5 +1,5 @@
 from collections.abc import Callable
-from typing import Any
+from typing import Any, cast
 
 import torch
 
@@ -20,7 +20,7 @@ def _tau(predictions: Tensor, targets: Tensor) -> float:
         np_preds = predictions.flatten().cpu().numpy()
         np_targets = targets.flatten().cpu().numpy()
         r = kendalltau(np_preds, np_targets, variant=variant).statistic
-        return r
+        return float(r)
 
     return _tau
 
@@ -121,4 +121,4 @@ def compute(self) -> float:
         if len(self._predictions) < 1 or len(self._targets) < 1:
             raise NotComputableError("KendallRankCorrelation must have at least one example before it can be computed.")
 
-        return super().compute()
+        return cast(float, super().compute())

ignite/metrics/regression/spearman_correlation.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, cast
 from collections.abc import Callable
 
 import torch
@@ -16,7 +16,7 @@ def _spearman_r(predictions: Tensor, targets: Tensor) -> float:
     np_preds = predictions.flatten().cpu().numpy()
     np_targets = targets.flatten().cpu().numpy()
     r = spearmanr(np_preds, np_targets).statistic
-    return r
+    return float(r)
 
 
 class SpearmanRankCorrelation(EpochMetric):
@@ -110,4 +110,4 @@ def compute(self) -> float:
                 "SpearmanRankCorrelation must have at least one example before it can be computed."
             )
 
-        return super().compute()
+        return cast(float, super().compute())