Repository: Lightning-AI/torchmetrics

Issue:
Nested MetricCollection doesn't add correct prefixes to Metric returning dict
## 🐛 Bug

When a metric returns dict, it doesn't correctly work with prefixes inside nested MetricCollection.

### To Reproduce

```
class MetricThatReturnsDict(Metric):
    def __init__(self, some_parameter):
        super().__init__()
        self.some_parameter = some_parameter
        self.add_state('data', default=[])

    def update(self, data):
        self.data.append(data)

    def compute(self):
        result = very_long_computation(self.data, some_parameter)
        return_dict = {
            'computed_result': result,
            'double_the_result': 2*result
        }
        return return_dict

collection = MetricCollection([
    MetricThatReturnsDict(42),
    ...
], prefix='first_')

parent_collection = MetricCollection([
    collection,
    ...
], prefix='second_')

parent_collection.update(tensor([1, 2, 3, 4]))
print(parent_collection.compute().keys())
>>> dict_keys(['second_computed_result', 'second_double_the_result'])
```
Expected result is
```
>>> dict_keys(['second_first_computed_result', 'second_first_double_the_result'])
```

If the metric returns a single value instead, the prefixes work correctly (`second_first_`).


### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 0.11.4, conda
- Python & PyTorch Version (e.g., 1.0): Python 3.9.16
- Any other relevant information such as OS (e.g., Linux): Linux

### Additional context

It could be possible to use two Metrics and return a single value from each one. But the `compute()` function would then be called two times separately, even though I could reuse the result of `very_long_computation`.


Discussion:
Hi! thanks for your contribution!, great first issue!

Resolving patch:
diff --git a/src/torchmetrics/collections.py b/src/torchmetrics/collections.py
--- a/src/torchmetrics/collections.py
+++ b/src/torchmetrics/collections.py
@@ -19,10 +19,11 @@
 import torch
 from torch import Tensor
 from torch.nn import Module, ModuleDict
+from typing_extensions import Literal
 
 from torchmetrics.metric import Metric
 from torchmetrics.utilities import rank_zero_warn
-from torchmetrics.utilities.data import _flatten_dict, allclose
+from torchmetrics.utilities.data import allclose
 from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
 from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val
 
@@ -176,9 +177,7 @@ def forward(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
         Positional arguments (args) will be passed to every metric in the collection, while keyword arguments (kwargs)
         will be filtered based on the signature of the individual metric.
         """
-        res = {k: m(*args, **m._filter_kwargs(**kwargs)) for k, m in self.items(keep_base=True, copy_state=False)}
-        res = _flatten_dict(res)
-        return {self._set_name(k): v for k, v in res.items()}
+        return self._compute_and_reduce("forward", *args, **kwargs)
 
     def update(self, *args: Any, **kwargs: Any) -> None:
         """Call update for each metric sequentially.
@@ -288,9 +287,43 @@ def _compute_groups_create_state_ref(self, copy: bool = False) -> None:
 
     def compute(self) -> Dict[str, Any]:
         """Compute the result for each metric in the collection."""
-        res = {k: m.compute() for k, m in self.items(keep_base=True, copy_state=False)}
-        res = _flatten_dict(res)
-        return {self._set_name(k): v for k, v in res.items()}
+        return self._compute_and_reduce("compute")
+
+    def _compute_and_reduce(
+        self, method_name: Literal["compute", "forward"], *args: Any, **kwargs: Any
+    ) -> Dict[str, Any]:
+        """Compute result from collection and reduce into a single dictionary.
+
+        Args:
+            method_name: The method to call on each metric in the collection.
+                Should be either `compute` or `forward`.
+            args: Positional arguments to pass to each metric (if method_name is `forward`)
+            kwargs: Keyword arguments to pass to each metric (if method_name is `forward`)
+
+        Raises:
+            ValueError:
+                If method_name is not `compute` or `forward`.
+
+        """
+        result = {}
+        for k, m in self.items(keep_base=True, copy_state=False):
+            if method_name == "compute":
+                res = m.compute()
+            elif method_name == "forward":
+                res = m(*args, **m._filter_kwargs(**kwargs))
+            else:
+                raise ValueError("method_name should be either 'compute' or 'forward', but got {method_name}")
+
+            if isinstance(res, dict):
+                for key, v in res.items():
+                    if hasattr(m, "prefix") and m.prefix is not None:
+                        key = f"{m.prefix}{key}"
+                    if hasattr(m, "postfix") and m.postfix is not None:
+                        key = f"{key}{m.postfix}"
+                    result[key] = v
+            else:
+                result[k] = res
+        return {self._set_name(k): v for k, v in result.items()}
 
     def reset(self) -> None:
         """Call reset for each metric sequentially."""
@@ -358,6 +391,8 @@ def add_metrics(
                     self[name] = metric
                 else:
                     for k, v in metric.items(keep_base=False):
+                        v.postfix = metric.postfix
+                        v.prefix = metric.prefix
                         self[f"{name}_{k}"] = v
         elif isinstance(metrics, Sequence):
             for metric in metrics:
@@ -373,9 +408,14 @@ def add_metrics(
                     self[name] = metric
                 else:
                     for k, v in metric.items(keep_base=False):
+                        v.postfix = metric.postfix
+                        v.prefix = metric.prefix
                         self[k] = v
         else:
-            raise ValueError("Unknown input to MetricCollection.")
+            raise ValueError(
+                "Unknown input to MetricCollection. Expected, `Metric`, `MetricCollection` or `dict`/`sequence` of the"
+                f" previous, but got {metrics}"
+            )
 
         self._groups_checked = False
         if self._enable_compute_groups:


Test patch:
diff --git a/tests/unittests/bases/test_collections.py b/tests/unittests/bases/test_collections.py
--- a/tests/unittests/bases/test_collections.py
+++ b/tests/unittests/bases/test_collections.py
@@ -36,7 +36,7 @@
 )
 from torchmetrics.utilities.checks import _allclose_recursive
 from unittests.helpers import seed_all
-from unittests.helpers.testers import DummyMetricDiff, DummyMetricSum
+from unittests.helpers.testers import DummyMetricDiff, DummyMetricMultiOutputDict, DummyMetricSum
 
 seed_all(42)
 
@@ -618,3 +618,13 @@ def test_nested_collections(input_collections):
     assert "valmetrics/macro_MulticlassPrecision" in val
     assert "valmetrics/micro_MulticlassAccuracy" in val
     assert "valmetrics/micro_MulticlassPrecision" in val
+
+
+def test_double_nested_collections():
+    """Test that double nested collections gets flattened to a single collection."""
+    collection1 = MetricCollection([DummyMetricMultiOutputDict()], prefix="prefix1_", postfix="_postfix1")
+    collection2 = MetricCollection([collection1], prefix="prefix2_", postfix="_postfix2")
+    x = torch.randn(10).sum()
+    val = collection2(x)
+    assert "prefix2_prefix1_output1_postfix1_postfix2" in val
+    assert "prefix2_prefix1_output2_postfix1_postfix2" in val
diff --git a/tests/unittests/helpers/testers.py b/tests/unittests/helpers/testers.py
--- a/tests/unittests/helpers/testers.py
+++ b/tests/unittests/helpers/testers.py
@@ -613,6 +613,14 @@ def compute(self):
         return [self.x, self.x]
 
 
+class DummyMetricMultiOutputDict(DummyMetricSum):
+    """DummyMetricMultiOutput for testing core components."""
+
+    def compute(self):
+        """Compute value."""
+        return {"output1": self.x, "output2": self.x}
+
+
 def inject_ignore_index(x: Tensor, ignore_index: int) -> Tensor:
     """Injecting the ignored index value into a tensor randomly."""
     if any(x.flatten() == ignore_index):  # ignore index is a class label