Fix corner case in MatthewsCorrcoef (#2743)

* fix + test * changelog * Apply suggestions from code review --------- Co-authored-by: Jirka Borovec <[email protected]>
Lightning-AI · Sep 16, 2024 · 32290ca · 32290ca
1 parent 62d9d32
commit 32290ca
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -63,6 +63,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed flakiness in tests related to `torch.unique` with `dim=None` ([#2650](https://github.com/Lightning-AI/torchmetrics/pull/2650))
 
 
+- Fixed corner case in `MatthewsCorrCoef` ([#2743](https://github.com/Lightning-AI/torchmetrics/pull/2743))
+
+
 ## [1.4.1] - 2024-08-02
 
 ### Changed

diff --git a/src/torchmetrics/functional/classification/matthews_corrcoef.py b/src/torchmetrics/functional/classification/matthews_corrcoef.py
@@ -64,12 +64,14 @@ def _matthews_corrcoef_reduce(confmat: Tensor) -> Tensor:
     denom = cov_ypyp * cov_ytyt
 
     if denom == 0 and confmat.numel() == 4:
-        if tp == 0 or tn == 0:
-            a = tp + tn
-
-        if fp == 0 or fn == 0:
-            b = fp + fn
-
+        if fn == 0 and tn == 0:
+            a, b = tp, fp
+        elif fp == 0 and tn == 0:
+            a, b = tp, fn
+        elif tp == 0 and fn == 0:
+            a, b = tn, fp
+        elif tp == 0 and fp == 0:
+            a, b = tn, fn
         eps = torch.tensor(torch.finfo(torch.float32).eps, dtype=torch.float32, device=confmat.device)
         numerator = torch.sqrt(eps) * (a - b)
         denom = (tp + fp + eps) * (tp + fn + eps) * (tn + fp + eps) * (tn + fn + eps)

diff --git a/tests/unittests/classification/test_matthews_corrcoef.py b/tests/unittests/classification/test_matthews_corrcoef.py
@@ -331,6 +331,12 @@ def test_zero_case_in_multiclass():
             torch.tensor([0, 0, 0, 0, 0, 1, 1, 1, 1, 1]),
             0.0,
         ),
+        (
+            binary_matthews_corrcoef,
+            torch.tensor([1, 1, 1, 1, 1, 0, 0, 0, 0, 0]),
+            torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0]),
+            0.0,
+        ),
         (binary_matthews_corrcoef, torch.zeros(10), torch.ones(10), -1.0),
         (binary_matthews_corrcoef, torch.ones(10), torch.zeros(10), -1.0),
         (