fix for MetricCollection.update gives identical results (#2944)

* fix issue 2916
* chlog
* linter

---------

Co-authored-by: Roman <Roman>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>
Co-authored-by: Jirka B <j.borovec+github@gmail.com>
Co-authored-by: Nicki Skafte Detlefsen <skaftenicki@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: 5 people

CHANGELOG.md

@@ -27,7 +27,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
--
+- Fixed `MetricCollection.update` gives identical results ([#2944](https://github.com/Lightning-AI/torchmetrics/issues/2944))
 
 
 ---

src/torchmetrics/collections.py

@@ -314,11 +314,19 @@ def _equal_metric_states(metric1: Metric, metric2: Metric) -> bool:
             if type(state1) != type(state2):  # noqa: E721
                 return False
 
-            if isinstance(state1, Tensor) and isinstance(state2, Tensor):
-                return state1.shape == state2.shape and allclose(state1, state2)
+            if (
+                isinstance(state1, Tensor)
+                and isinstance(state2, Tensor)
+                and not (state1.shape == state2.shape and allclose(state1, state2))
+            ):
+                return False
 
-            if isinstance(state1, list) and isinstance(state2, list):
-                return all(s1.shape == s2.shape and allclose(s1, s2) for s1, s2 in zip(state1, state2))
+            if (
+                isinstance(state1, list)
+                and isinstance(state2, list)
+                and not (all(s1.shape == s2.shape and allclose(s1, s2) for s1, s2 in zip(state1, state2)))
+            ):
+                return False
 
         return True
 

tests/unittests/bases/test_collections.py

@@ -33,6 +33,7 @@
     MultilabelAUROC,
     MultilabelAveragePrecision,
 )
+from torchmetrics.text import BLEUScore
 from torchmetrics.utilities.checks import _allclose_recursive
 from unittests._helpers import seed_all
 from unittests._helpers.testers import DummyMetricDiff, DummyMetricMultiOutputDict, DummyMetricSum
@@ -743,3 +744,33 @@ def compute(self):
     # Print the calculated metrics
     assert "my_prefix/accuracy/my_postfix" in res
     assert "my_prefix/precision/my_postfix" in res
+
+
+def test_collection_update():
+    """Test that metric collection updates metrics.
+
+    See issue: https://github.com/Lightning-AI/torchmetrics/issues/2916
+
+    """
+    metrics = MetricCollection({
+        "bleu-1": BLEUScore(1),
+        "bleu-2": BLEUScore(2),
+        "bleu-3": BLEUScore(3),
+        "bleu-4": BLEUScore(4),
+    })
+
+    preds = ["the cat is on the mat"]
+    target = [["there is a cat on the mat", "a cat is on the mat"]]
+
+    metrics.update(preds, target)
+    actual = metrics.compute()
+
+    expected = {
+        "bleu-1": torch.tensor(0.8333),
+        "bleu-2": torch.tensor(0.8165),
+        "bleu-3": torch.tensor(0.7937),
+        "bleu-4": torch.tensor(0.7598),
+    }
+
+    for k, v in expected.items():
+        torch.testing.assert_close(actual=actual.get(k), expected=v, rtol=1e-4, atol=1e-4)