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
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import typing
 from typing import Any, Dict, List, Optional, Sequence, Union
 
 from torch import Tensor
@@ -20,6 +21,9 @@
 from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
 from torchmetrics.wrappers.abstract import WrapperMetric
 
+if typing.TYPE_CHECKING:
+    from torch.nn import Module
+
 if not _MATPLOTLIB_AVAILABLE:
     __doctest_skip__ = ["ClasswiseWrapper.plot"]
 
@@ -209,3 +213,22 @@ def plot(
 
         """
         return self._plot(val, ax)
+
+    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)
+
+        return super().__getattr__(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.
+            setattr(self.metric, name, 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
@@ -17,7 +17,7 @@
 
 import pytest
 import torch
-from torchmetrics import Metric, MetricCollection
+from torchmetrics import ClasswiseWrapper, Metric, MetricCollection
 from torchmetrics.classification import (
     BinaryAccuracy,
     MulticlassAccuracy,
@@ -540,6 +540,44 @@ def test_compute_group_define_by_user():
     assert m.compute()
 
 
+def test_classwise_wrapper_compute_group():
+    """Check that user can provide compute groups."""
+    classwise_accuracy = ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None), prefix="accuracy")
+    classwise_recall = ClasswiseWrapper(MulticlassRecall(num_classes=3, average=None), prefix="recall")
+    classwise_precision = ClasswiseWrapper(MulticlassPrecision(num_classes=3, average=None), prefix="precision")
+
+    m = MetricCollection(
+        {
+            "accuracy": ClasswiseWrapper(MulticlassAccuracy(num_classes=3, average=None), prefix="accuracy"),
+            "recall": ClasswiseWrapper(MulticlassRecall(num_classes=3, average=None), prefix="recall"),
+            "precision": ClasswiseWrapper(MulticlassPrecision(num_classes=3, average=None), prefix="precision"),
+        },
+        compute_groups=[["accuracy", "recall", "precision"]],
+    )
+
+    # Check that we are not going to check the groups in the first update
+    assert m._groups_checked
+    assert m.compute_groups == {0: ["accuracy", "recall", "precision"]}
+
+    preds = torch.randn(10, 3).softmax(dim=-1)
+    target = torch.randint(3, (10,))
+
+    expected = {
+        **classwise_accuracy(preds, target),
+        **classwise_recall(preds, target),
+        **classwise_precision(preds, target),
+    }
+
+    m.update(preds, target)
+    res = m.compute()
+
+    for key in expected:
+        assert torch.allclose(res[key], expected[key])
+
+    # check metric state_dict
+    m.state_dict()
+
+
 def test_compute_on_different_dtype():
     """Check that extraction of compute groups are robust towards difference in dtype."""
     m = MetricCollection([