fix logic + tests

Author: SkafteNicki

src/torchmetrics/collections.py

@@ -247,10 +247,7 @@ 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))
-            if self._state_is_copy:
-                # If we have deep copied state in between updates, reestablish link
-                self._compute_groups_create_state_ref()
-                self._state_is_copy = False
+            self._compute_groups_create_state_ref()
         else:  # the first update always do per metric to form compute groups
             for m in self.values(copy_state=False):
                 m_kwargs = m._filter_kwargs(**kwargs)
@@ -339,16 +336,15 @@ def _compute_groups_create_state_ref(self, copy: bool = False) -> None:
                 of just passed by reference
 
         """
-        if not self._state_is_copy and self._groups_checked:
-            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
-                        setattr(mi, state, deepcopy(m0_state) if copy else m0_state)
-                    mi._update_count = deepcopy(m0._update_count) if copy else m0._update_count
+        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
+                    setattr(mi, state, deepcopy(m0_state) if copy else m0_state)
+                mi._update_count = deepcopy(m0._update_count) if copy else m0._update_count
         self._state_is_copy = copy
 
     def compute(self) -> dict[str, Any]:

tests/unittests/bases/test_collections.py

@@ -33,6 +33,7 @@
     MultilabelAUROC,
     MultilabelAveragePrecision,
 )
+from torchmetrics.regression import PearsonCorrCoef
 from torchmetrics.text import BLEUScore
 from torchmetrics.utilities.checks import _allclose_recursive
 from unittests._helpers import seed_all
@@ -796,3 +797,39 @@ def test_collection_update():
 
     for k, v in expected.items():
         torch.testing.assert_close(actual=actual.get(k), expected=v, rtol=1e-4, atol=1e-4)
+
+
+def test_collection_state_being_re_established_after_copy():
+    """Check that shared metrics states when using compute groups are re-established after a copy.
+
+    See issue: https://github.com/Lightning-AI/torchmetrics/issues/2896
+
+    """
+    m1, m2 = PearsonCorrCoef(), PearsonCorrCoef()
+    m12 = MetricCollection({"m1": m1, "m2": m2}, compute_groups=True)
+    x1, y1 = torch.randn(100), torch.randn(100)
+    m12.update(x1, y1)
+    assert m12.compute_groups == {0: ["m1", "m2"]}
+
+    # Check that the states are pointing to the same location
+    assert not m12._state_is_copy
+    assert m12.m1.mean_x.data_ptr() == m12.m2.mean_x.data_ptr(), "States should point to the same location"
+
+    # Break the references between the states
+    _ = m12.items()
+    assert m12._state_is_copy
+    assert m12.m1.mean_x.data_ptr() != m12.m2.mean_x.data_ptr(), "States should not point to the same location"
+
+    # Update should restore the references between the states
+    x2, y2 = torch.randn(100), torch.randn(100)
+
+    m12.update(x2, y2)
+    assert not m12._state_is_copy
+    assert m12.m1.mean_x.data_ptr() == m12.m2.mean_x.data_ptr(), "States should point to the same location"
+
+    x3, y3 = torch.randn(100), torch.randn(100)
+    m12.update(x3, y3)
+
+    assert not m12._state_is_copy
+    assert m12.m1.mean_x.data_ptr() == m12.m2.mean_x.data_ptr(), "States should point to the same location"
+    assert m12._equal_metric_states(m12.m1, m12.m2)