Repository: Lightning-AI/torchmetrics

Issue:
MetricCollection doesn't merge compute groups on forward
## 🐛 Bug

So I'm not entirely sure whether this is a bug or intentional, in that case this would be a doc fix request.

As the title states, the `MetricCollection` will not merge compute groups on `forward()`. While `update()` merges the compute groups and calls update on all metrics, forward will only call forward on all metrics and ignore compute groups altogether.

This is contradictory to the example in the [docs](https://torchmetrics.readthedocs.io/en/stable/pages/overview.html#metriccollection), which shows an example using forward and claims this procedure would gain a speedup from compute groups.

### To Reproduce

Setup a `MetricCollection` with merge-able Metrics, call forward on it and observe no merging:

#### Code sample

```python
import torch, torchmetrics

m = torchmetrics.MetricCollection(torchmetrics.Accuracy(), torchmetrics.Recall())
# Binary inputs
binary_preds  = torch.tensor([0.6, 0.1, 0.9])
binary_target = torch.tensor([1, 0, 1])

m._groups
# Out[28]: {0: ['Accuracy'], 1: ['Recall']} - initially not merged, this is fine

m(binary_preds, binary_target)
# Out[31]: {'Accuracy': tensor(1.), 'Recall': tensor(1.)}
m._groups
# Out[32]: {0: ['Accuracy'], 1: ['Recall']} - still not merged, this is bad(?)
m.update(binary_preds, binary_target)
m._groups
# Out[34]: {0: ['Accuracy', 'Recall']} - update correctly merges the metrics
```

### Expected behavior

Merge Groups should be merged and used, I think, i.e. forward should behave as update in that regard. Alternatively I can see how, since forward and update behave differently (one returning the metric, the other not), compute groups might not be as effective here. In that case, the docs should be updated to not promise compute group speedups on forward calls, and ideally a warning describing this behavior should be added.

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): `0.9.1`, `pip`
- Python & PyTorch Version (e.g., 1.0):  `3.8.12` & `1.11.0`
- Any other relevant information such as OS (e.g., Linux): `FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel`



Discussion:

Resolving patch:
diff --git a/src/torchmetrics/collections.py b/src/torchmetrics/collections.py
--- a/src/torchmetrics/collections.py
+++ b/src/torchmetrics/collections.py
@@ -54,6 +54,14 @@ class name as key for the output dict.
             this argument is ``True`` which enables this feature. Set this argument to `False` for disabling
             this behaviour. Can also be set to a list of lists of metrics for setting the compute groups yourself.
 
+    .. note::
+        The compute groups feature can significatly speedup the calculation of metrics under the right conditions.
+        First, the feature is only available when calling the ``update`` method and not when calling ``forward`` method
+        due to the internal logic of ``forward`` preventing this. Secondly, since we compute groups share metric
+        states by reference, calling ``.items()``, ``.values()`` etc. on the metric collection will break this
+        reference and a copy of states are instead returned in this case (reference will be reestablished on the next
+        call to ``update``).
+
     .. note::
         Metric collections can be nested at initilization (see last example) but the output of the collection will
         still be a single flatten dictionary combining the prefix and postfix arguments from the nested collection.
@@ -101,13 +109,16 @@ class name as key for the output dict.
 
     Example (specification of compute groups):
         >>> metrics = MetricCollection(
-        ...     Accuracy(),
+        ...     Recall(num_classes=3, average='macro'),
         ...     Precision(num_classes=3, average='macro'),
         ...     MeanSquaredError(),
-        ...     compute_groups=[['Accuracy', 'Precision'], ['MeanSquaredError']]
+        ...     compute_groups=[['Recall', 'Precision'], ['MeanSquaredError']]
         ... )
-        >>> pprint(metrics(preds, target))
-        {'Accuracy': tensor(0.1250), 'MeanSquaredError': tensor(2.3750), 'Precision': tensor(0.0667)}
+        >>> metrics.update(preds, target)
+        >>> pprint(metrics.compute())
+        {'MeanSquaredError': tensor(2.3750), 'Precision': tensor(0.0667), 'Recall': tensor(0.1111)}
+        >>> pprint(metrics.compute_groups)
+        {0: ['Recall', 'Precision'], 1: ['MeanSquaredError']}
 
     Example (nested metric collections):
         >>> metrics = MetricCollection([
@@ -170,9 +181,6 @@ def update(self, *args: Any, **kwargs: Any) -> None:
                 # only update the first member
                 m0 = getattr(self, cg[0])
                 m0.update(*args, **m0._filter_kwargs(**kwargs))
-                for i in range(1, len(cg)):  # copy over the update count
-                    mi = getattr(self, cg[i])
-                    mi._update_count = m0._update_count
             if self._state_is_copy:
                 # If we have deep copied state inbetween updates, reestablish link
                 self._compute_groups_create_state_ref()
@@ -191,7 +199,8 @@ def update(self, *args: Any, **kwargs: Any) -> None:
     def _merge_compute_groups(self) -> None:
         """Iterates over the collection of metrics, checking if the state of each metric matches another.
 
-        If so, their compute groups will be merged into one
+        If so, their compute groups will be merged into one. The complexity of the method is approximately
+        ``O(number_of_metrics_in_collection ** 2)``, as all metrics need to be compared to all other metrics.
         """
         n_groups = len(self._groups)
         while True:
@@ -264,6 +273,7 @@ def _compute_groups_create_state_ref(self, copy: bool = False) -> None:
                         m0_state = getattr(m0, state)
                         # Determine if we just should set a reference or a full copy
                         setattr(mi, state, deepcopy(m0_state) if copy else m0_state)
+                    setattr(mi, "_update_count", deepcopy(m0._update_count) if copy else m0._update_count)
         self._state_is_copy = copy
 
     def compute(self) -> Dict[str, Any]:


Test patch: