Repository: Lightning-AI/torchmetrics

Issue:
Incorrect caching (`m._compute`) of metrics inside a `MetricCollection` if `compute_groups` are used and `.compute` is called twice
## 🐛 Bug

Hello Team! I have found a bug which I already tracked down to is related to the caching attribute `_compute` of a metric in a metric collection. If the `.compute()` method is called twice then the second call returns different (and false) results if compute_groups are used. That is because the first `.compute()` call will store result of the first metric in the `_compute` cache for all other metrics of that compute group, which is then returned in the second call.

### To Reproduce

Steps to reproduce the behavior...

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

<details>
  <summary>Code sample</summary>

```py
import torch
from torchmetrics import AUROC, ROC, Recall, F1Score, MetricCollection

metrics = MetricCollection({
    "auroc": AUROC(task='binary'),
    "roc": ROC(task='binary'),
    "recall": Recall(task='binary'),
    "f1": F1Score(task='binary')
})

y_true = torch.tensor([1, 0, 0, 1])
y_pred = torch.tensor([0.6, 0.2, 0.4, 0.2])

for batch in range(10):
    metrics.update(y_pred, y_true)

print("First compute call:")
print(metrics.compute())
print("Second (and all consecutive) compute call:")
print(metrics.compute())
print("Compute Groups:")
print(metrics.compute_groups)
```

This results in:

```sh
First compute call:
{
    'auroc': tensor(0.6250),
    'f1': tensor(0.6667),
    'recall': tensor(0.5000),
    'roc': (
        tensor([0.0000, 0.0000, 0.5000, 1.0000]),
        tensor([0.0000, 0.5000, 0.5000, 1.0000]),
        tensor([1.0000, 0.6000, 0.4000, 0.2000])
    )
}
Second (and all consecutive) compute call:
{'auroc': tensor(0.6250), 'f1': tensor(0.6667), 'recall': tensor(0.6667), 'roc': tensor(0.6250)}
Compute Groups:
{0: ['auroc', 'roc'], 1: ['f1', 'recall']}
```

<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

</details>

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

<details>
  <summary>I expect the output of above code to be</summary>

```sh
First compute call:
{
    'auroc': tensor(0.6250),
    'f1': tensor(0.6667),
    'recall': tensor(0.5000),
    'roc': (
        tensor([0.0000, 0.0000, 0.5000, 1.0000]),
        tensor([0.0000, 0.5000, 0.5000, 1.0000]),
        tensor([1.0000, 0.6000, 0.4000, 0.2000])
    )
}
Second (and all consecutive) compute call:
{
    'auroc': tensor(0.6250),
    'f1': tensor(0.6667),
    'recall': tensor(0.5000),
    'roc': (
        tensor([0.0000, 0.0000, 0.5000, 1.0000]),
        tensor([0.0000, 0.5000, 0.5000, 1.0000]),
        tensor([1.0000, 0.6000, 0.4000, 0.2000])
    )
}
Compute Groups:
{0: ['auroc', 'roc'], 1: ['f1', 'recall']}
```
</details>

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source):  1.4.0.post0 installed via `rye add`
- Python & PyTorch Version (e.g., 1.0): Python 3.10.14 & PyTorch 2.2.0
- Any other relevant information such as OS (e.g., Linux): Ubuntu 22 on WSL

### Additional context

I already tracked down the respective row:

https://github.com/Lightning-AI/torchmetrics/blob/8d9d2ae5f6fea65e861026b90e365ec34b5c57b1/src/torchmetrics/collections.py#L307



It was introduced by the PR #2211 . Since they also fixed something with that, simply reverting the change would not help. I currently try to find an other solution for this.

_But Tobi, why do you want to have two compute calls in the first place?_

The reason why I want to use two seperatly `.compute()` calls is that in my current code structure I want to log my metrics in two different functions which have access to the same metrics object. The use case is actually pretty close to the example I provided above. Of course I could use a workaround by just storing the result by myself and handle it in my respective functions, but this would make my code less readable. Therefore a fix for this Issue would be a nice to have for me. **However, other people could run into the same problem without noticing, since for some e.g. scalar metrics the code would run as usual, just with the wrong outputs.**




<!-- Add any other context about the problem here. -->



Discussion:
1. I found out that test for #2211 are not properly implemented: Deleting the added row from that PR also results in all tests pass, validating it manually, however, shows that the bug #2211 fixed exists if the respective row is not present.
2. I also found out that some contents of compute could be references instead of deepcopies, which results in even weirder side effects.

I fixed all 3 things (the original bug and the two new findings). I will open a PR for this in a sec.

Resolving patch:
diff --git a/src/torchmetrics/collections.py b/src/torchmetrics/collections.py
--- a/src/torchmetrics/collections.py
+++ b/src/torchmetrics/collections.py
@@ -187,6 +187,11 @@ def __init__(
 
         self.add_metrics(metrics, *additional_metrics)
 
+    @property
+    def metric_state(self) -> Dict[str, Dict[str, Any]]:
+        """Get the current state of the metric."""
+        return {k: m.metric_state for k, m in self.items(keep_base=False, copy_state=False)}
+
     @torch.jit.unused
     def forward(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
         """Call forward for each metric sequentially.
@@ -206,6 +211,11 @@ def update(self, *args: Any, **kwargs: Any) -> None:
         """
         # Use compute groups if already initialized and checked
         if self._groups_checked:
+            # Delete the cache of all metrics to invalidate the cache and therefore recent compute calls, forcing new
+            # compute calls to recompute
+            for k in self.keys(keep_base=True):
+                mi = getattr(self, str(k))
+                mi._computed = None
             for cg in self._groups.values():
                 # only update the first member
                 m0 = getattr(self, cg[0])
@@ -304,7 +314,6 @@ def _compute_groups_create_state_ref(self, copy: bool = False) -> None:
                         # Determine if we just should set a reference or a full copy
                         setattr(mi, state, deepcopy(m0_state) if copy else m0_state)
                     mi._update_count = deepcopy(m0._update_count) if copy else m0._update_count
-                    mi._computed = deepcopy(m0._computed) if copy else m0._computed
         self._state_is_copy = copy
 
     def compute(self) -> Dict[str, Any]:
diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py
--- a/src/torchmetrics/metric.py
+++ b/src/torchmetrics/metric.py
@@ -630,6 +630,8 @@ def wrapped_func(*args: Any, **kwargs: Any) -> Any:
                 should_unsync=self._should_unsync,
             ):
                 value = _squeeze_if_scalar(compute(*args, **kwargs))
+                # clone tensor to avoid in-place operations after compute, altering already computed results
+                value = apply_to_collection(value, Tensor, lambda x: x.clone())
 
             if self.compute_with_cache:
                 self._computed = value


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
@@ -452,10 +452,29 @@ def test_check_compute_groups_correctness(self, metrics, expected, preds, target
             for key in res_cg:
                 assert torch.allclose(res_cg[key], res_without_cg[key])
 
+            # Check if second compute is the same
+            res_cg2 = m.compute()
+            for key in res_cg2:
+                assert torch.allclose(res_cg[key], res_cg2[key])
+
             if with_reset:
                 m.reset()
                 m2.reset()
 
+        # Test if a second compute without a reset is the same
+        m.reset()
+        m.update(preds, target)
+        res_cg = m.compute()
+        # Simulate different preds by simply inversing them
+        m.update(1 - preds, target)
+        res_cg2 = m.compute()
+        # Now check if the results from the first compute are different from the second
+        for key in res_cg:
+            # A different shape is okay, therefore skip (this happens for multidim_average="samplewise")
+            if res_cg[key].shape != res_cg2[key].shape:
+                continue
+            assert not torch.all(res_cg[key] == res_cg2[key])
+
     @pytest.mark.parametrize("method", ["items", "values", "keys"])
     def test_check_compute_groups_items_and_values(self, metrics, expected, preds, target, method):
         """Check states are copied instead of passed by ref when a single metric in the collection is access."""