Repository: Lightning-AI/torchmetrics

Issue:
`MetricCollection` did not copy inner state of metric in `ClasswiseWrapper` when computing groups metrics
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

`MetricCollection` did not copy inner state of metric in `ClasswiseWrapper` when computing groups metrics.

### 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>

<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

```python
import torch
from lightning import seed_everything
from torchmetrics import MetricCollection, ClasswiseWrapper
from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score

seed_everything(42)

random_pred = torch.rand((10, 3))
pred = torch.softmax(random_pred, dim=-1)
pred_class = torch.argmax(pred, dim=-1)
target = torch.randint(0, 3, size=(10,))

multiclass_acc = MulticlassAccuracy(
    num_classes=3,
    average=None,
)
print("multiclass_acc:", multiclass_acc(pred, target))

multiclass_f1 = MulticlassF1Score(
    num_classes=3,
    average=None,
)
print("multiclass_f1:", multiclass_f1(pred, target))

mc = MetricCollection(
    {
        "accuracy": ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None)),
        "f1": ClasswiseWrapper(MulticlassF1Score(num_classes=3, average=None)),
    },
    compute_groups=[
        ["accuracy", "f1"],
    ],
)

print("MetricCollection.forward:", mc(pred, target))
mc.reset()
mc.update(pred, target)
print("MetricCollection.update&compute:", mc.compute())
```

</details>

<details>
  <summary>Output</summary>

<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

```console
Seed set to 42
site-packages\torchmetrics\utilities\prints.py:43: UserWarning: The ``compute`` method of metric MulticlassF1Score was called before the ``update`` method which may lead to errors, as metric states have not yet been updated.
  warnings.warn(*args, **kwargs)  # noqa: B028
multiclass_acc: tensor([0.3333, 0.5000, 0.3333])
multiclass_f1: tensor([0.4000, 0.4444, 0.3333])
MetricCollection.forward: {'multiclassaccuracy_0': tensor(0.3333), 'multiclassaccuracy_1': tensor(0.5000), 'multiclassaccuracy_2': tensor(0.3333), 'multiclassf1score_0': tensor(0.4000), 'multiclassf1score_1': tensor(0.4444), 'multiclassf1score_2': tensor(0.3333)}
MetricCollection.update&compute: {'multiclassaccuracy_0': tensor(0.3333), 'multiclassaccuracy_1': tensor(0.5000), 'multiclassaccuracy_2': tensor(0.3333), 'multiclassf1score_0': tensor(0.), 'multiclassf1score_1': tensor(0.), 'multiclassf1score_2': tensor(0.)}
```

</details>

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

The metrics `multiclassf1score_0`, `multiclassf1score_1`, and `multiclassf1score_2` of `MetricCollection.compute` should be the same as `MetricCollection.forward` or the result of each metric.




<details>
  <summary>Solution 1</summary>

https://github.com/Lightning-AI/torchmetrics/blob/master/src/torchmetrics/collections.py#L305

```python
class MetricCollection:
    def _compute_groups_create_state_ref(self, copy: bool = False) -> None:
        """Create reference between metrics in the same compute group.

        Args:
            copy: If `True` the metric state will between members will be copied instead
                of just passed by reference

        """
        if not self._state_is_copy:
            for cg in self._groups.values():
                m0 = getattr(self, cg[0])
                for i in range(1, len(cg)):
                    mi = getattr(self, cg[i])
                    for state in m0._defaults:
                        m0_state = getattr(m0, state)
                        # Determine if we just should set a reference or a full copy
                        if isinstance(mi, ClasswiseWrapper):  # << Added
                            setattr(mi.metric, state, deepcopy(m0_state) if copy else m0_state) # << Added
                        setattr(mi, state, deepcopy(m0_state) if copy else m0_state)

                    if isinstance(mi, ClasswiseWrapper): # << Added
                        mi.metric._update_count = deepcopy(m0._update_count) if copy else m0._update_count # << Added
                        mi.metric._computed = deepcopy(m0._computed) if copy else m0._computed # << Added
                    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

```
</details>


<details>
  <summary>Solution 2</summary>

https://github.com/Lightning-AI/torchmetrics/blob/master/src/torchmetrics/wrappers/classwise.py#L27

```python
class ClasswiseWrapper:
    def __getattr__(self, name: str):
        # return state from self.metric
        if name in ["tp", "fp", "fn", "tn"]:   # <<Added
            return getattr(self.metric, name)  # <<Added

        return super().__getattr__(name)

    def __setattr__(self, name: str, value: Any) -> None:
        if hasattr(self, "metric") and name in ["tp", "fp", "fn", "tn"]:  # <<Added
            setattr(self.metric, name, value)                             # <<Added
        else:                                                             # <<Added
            super().__setattr__(name, value)
            if name == "metric":                                          # <<Added
                self._defaults = self.metric._defaults                    # <<Added


```
</details>



### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): 1.3.1
- Python & PyTorch Version (e.g., 1.0):
  - Python: 3.8
  - Pytorch: 2.1.1
- Any other relevant information such as OS (e.g., Linux): Windows 11

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Update Solution 2, which overrides the `__getattr__` and `__setattr__` of ClasswiseWrapper

Resolving patch:
diff --git a/src/torchmetrics/wrappers/classwise.py b/src/torchmetrics/wrappers/classwise.py
--- a/src/torchmetrics/wrappers/classwise.py
+++ b/src/torchmetrics/wrappers/classwise.py
@@ -216,19 +216,19 @@ def plot(
 
     def __getattr__(self, name: str) -> Union[Tensor, "Module"]:
         """Get attribute from classwise wrapper."""
-        # return state from self.metric
-        if name in ["tp", "fp", "fn", "tn"]:
-            return getattr(self.metric, name)
+        if name == "metric" or (name in self.__dict__ and name not in self.metric.__dict__):
+            # we need this to prevent from infinite getattribute loop.
+            return super().__getattr__(name)
 
-        return super().__getattr__(name)
+        return getattr(self.metric, name)
 
     def __setattr__(self, name: str, value: Any) -> None:
         """Set attribute to classwise wrapper."""
-        super().__setattr__(name, value)
-        if name == "metric":
-            self._defaults = self.metric._defaults
-            self._persistent = self.metric._persistent
-            self._reductions = self.metric._reductions
-        if hasattr(self, "metric") and name in ["tp", "fp", "fn", "tn", "_update_count", "_computed"]:
-            # update ``_update_count`` and ``_computed`` of internal metric to prevent warning.
+        if hasattr(self, "metric") and name in self.metric._defaults:
             setattr(self.metric, name, value)
+        else:
+            super().__setattr__(name, value)
+            if name == "metric":
+                self._defaults = self.metric._defaults
+                self._persistent = self.metric._persistent
+                self._reductions = self.metric._reductions


Test patch: