Merge branch 'master' into lint/ruff

Author: Borda

CHANGELOG.md

@@ -36,9 +36,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Fixed dtype being changed by deepspeed for certain regression metrics ([#2379](https://github.com/Lightning-AI/torchmetrics/pull/2379))
 
 
+- Fixed plotting of metric collection when prefix/postfix is set ([#2429](https://github.com/Lightning-AI/torchmetrics/pull/2429))
+
+
 - Fixed bug when `top_k>1` and `average="macro"` for classification metrics ([#2423](https://github.com/Lightning-AI/torchmetrics/pull/2423))
 
 
+- Fixed case where label prediction tensors in classification metrics were not validated correctly ([#2427](https://github.com/Lightning-AI/torchmetrics/pull/2427))
+
+
 - Fixed how auc scores are calculated in `PrecisionRecallCurve.plot` methods ([#2437](https://github.com/Lightning-AI/torchmetrics/pull/2437))
 
 ## [1.3.1] - 2024-02-12

src/torchmetrics/collections.py

@@ -647,12 +647,11 @@ def plot(
                     f"Expected argument `ax` to be a sequence of matplotlib axis objects with the same length as the "
                     f"number of metrics in the collection, but got {type(ax)} with len {len(ax)} when `together=False`"
                 )
-
         val = val or self.compute()
         if together:
             return plot_single_or_multi_val(val, ax=ax)
         fig_axs = []
-        for i, (k, m) in enumerate(self.items(keep_base=True, copy_state=False)):
+        for i, (k, m) in enumerate(self.items(keep_base=False, copy_state=False)):
             if isinstance(val, dict):
                 f, a = m.plot(val[k], ax=ax[i] if ax is not None else ax)
             elif isinstance(val, Sequence):

src/torchmetrics/functional/classification/confusion_matrix.py

@@ -285,21 +285,13 @@ def _multiclass_confusion_matrix_tensor_validation(
             " and `preds` should be (N, C, ...)."
         )
 
-    num_unique_values = len(torch.unique(target))
-    check = num_unique_values > num_classes if ignore_index is None else num_unique_values > num_classes + 1
-    if check:
-        raise RuntimeError(
-            "Detected more unique values in `target` than `num_classes`. Expected only "
-            f"{num_classes if ignore_index is None else num_classes + 1} but found "
-            f"{num_unique_values} in `target`."
-        )
-
-    if not preds.is_floating_point():
-        num_unique_values = len(torch.unique(preds))
-        if num_unique_values > num_classes:
+    check_value = num_classes if ignore_index is None else num_classes + 1
+    for t, name in ((target, "target"),) + ((preds, "preds"),) if not preds.is_floating_point() else ():  # noqa: RUF005
+        num_unique_values = len(torch.unique(t))
+        if num_unique_values > check_value:
             raise RuntimeError(
-                "Detected more unique values in `preds` than `num_classes`. Expected only "
-                f"{num_classes} but found {num_unique_values} in `preds`."
+                f"Detected more unique values in `{name}` than expected. Expected only {check_value} but found"
+                f" {num_unique_values} in `target`."
             )
 
 

src/torchmetrics/functional/classification/stat_scores.py

@@ -304,21 +304,13 @@ def _multiclass_stat_scores_tensor_validation(
             " and `preds` should be (N, C, ...)."
         )
 
-    num_unique_values = len(torch.unique(target))
-    check = num_unique_values > num_classes if ignore_index is None else num_unique_values > num_classes + 1
-    if check:
-        raise RuntimeError(
-            "Detected more unique values in `target` than `num_classes`. Expected only"
-            f" {num_classes if ignore_index is None else num_classes + 1} but found"
-            f" {num_unique_values} in `target`."
-        )
-
-    if not preds.is_floating_point():
-        unique_values = torch.unique(preds)
-        if len(unique_values) > num_classes:
+    check_value = num_classes if ignore_index is None else num_classes + 1
+    for t, name in ((target, "target"),) + ((preds, "preds"),) if not preds.is_floating_point() else ():  # noqa: RUF005
+        num_unique_values = len(torch.unique(t))
+        if num_unique_values > check_value:
             raise RuntimeError(
-                "Detected more unique values in `preds` than `num_classes`. Expected only"
-                f" {num_classes} but found {len(unique_values)} in `preds`."
+                f"Detected more unique values in `{name}` than expected. Expected only {check_value} but found"
+                f" {num_unique_values} in `target`."
             )
 
 

tests/unittests/classification/test_confusion_matrix.py

@@ -239,6 +239,41 @@ def test_multiclass_confusion_matrix_dtype_gpu(self, inputs, dtype):
         )
 
 
+@pytest.mark.parametrize(
+    ("preds", "target", "ignore_index", "error_message"),
+    [
+        (
+            torch.randint(NUM_CLASSES + 1, (100,)),
+            torch.randint(NUM_CLASSES, (100,)),
+            None,
+            f"Detected more unique values in `preds` than expected. Expected only {NUM_CLASSES}.*",
+        ),
+        (
+            torch.randint(NUM_CLASSES, (100,)),
+            torch.randint(NUM_CLASSES + 1, (100,)),
+            None,
+            f"Detected more unique values in `target` than expected. Expected only {NUM_CLASSES}.*",
+        ),
+        (
+            torch.randint(NUM_CLASSES + 2, (100,)),
+            torch.randint(NUM_CLASSES, (100,)),
+            1,
+            f"Detected more unique values in `preds` than expected. Expected only {NUM_CLASSES + 1}.*",
+        ),
+        (
+            torch.randint(NUM_CLASSES, (100,)),
+            torch.randint(NUM_CLASSES + 2, (100,)),
+            1,
+            f"Detected more unique values in `target` than expected. Expected only {NUM_CLASSES + 1}.*",
+        ),
+    ],
+)
+def test_raises_error_on_too_many_classes(preds, target, ignore_index, error_message):
+    """Test that an error is raised if the number of classes in preds or target is larger than expected."""
+    with pytest.raises(RuntimeError, match=error_message):
+        multiclass_confusion_matrix(preds, target, num_classes=NUM_CLASSES, ignore_index=ignore_index)
+
+
 def test_multiclass_overflow():
     """Test that multiclass computations does not overflow even on byte inputs."""
     preds = torch.randint(20, (100,)).byte()

tests/unittests/classification/test_stat_scores.py

@@ -325,6 +325,41 @@ def test_multiclass_stat_scores_dtype_gpu(self, inputs, dtype):
         )
 
 
+@pytest.mark.parametrize(
+    ("preds", "target", "ignore_index", "error_message"),
+    [
+        (
+            torch.randint(NUM_CLASSES + 1, (100,)),
+            torch.randint(NUM_CLASSES, (100,)),
+            None,
+            f"Detected more unique values in `preds` than expected. Expected only {NUM_CLASSES}.*",
+        ),
+        (
+            torch.randint(NUM_CLASSES, (100,)),
+            torch.randint(NUM_CLASSES + 1, (100,)),
+            None,
+            f"Detected more unique values in `target` than expected. Expected only {NUM_CLASSES}.*",
+        ),
+        (
+            torch.randint(NUM_CLASSES + 2, (100,)),
+            torch.randint(NUM_CLASSES, (100,)),
+            1,
+            f"Detected more unique values in `preds` than expected. Expected only {NUM_CLASSES + 1}.*",
+        ),
+        (
+            torch.randint(NUM_CLASSES, (100,)),
+            torch.randint(NUM_CLASSES + 2, (100,)),
+            1,
+            f"Detected more unique values in `target` than expected. Expected only {NUM_CLASSES + 1}.*",
+        ),
+    ],
+)
+def test_raises_error_on_too_many_classes(preds, target, ignore_index, error_message):
+    """Test that an error is raised if the number of classes in preds or target is larger than expected."""
+    with pytest.raises(RuntimeError, match=error_message):
+        multiclass_stat_scores(preds, target, num_classes=NUM_CLASSES, ignore_index=ignore_index)
+
+
 _mc_k_target = torch.tensor([0, 1, 2])
 _mc_k_preds = torch.tensor([[0.35, 0.4, 0.25], [0.1, 0.5, 0.4], [0.2, 0.1, 0.7]])
 

tests/unittests/utilities/test_plot.py

@@ -834,12 +834,17 @@ def test_confusion_matrix_plotter(metric_class, preds, target, labels, use_label
 
 @pytest.mark.parametrize("together", [True, False])
 @pytest.mark.parametrize("num_vals", [1, 2])
-def test_plot_method_collection(together, num_vals):
+@pytest.mark.parametrize(
+    ("prefix", "postfix"), [(None, None), ("prefix", None), (None, "postfix"), ("prefix", "postfix")]
+)
+def test_plot_method_collection(together, num_vals, prefix, postfix):
     """Test the plot method of metric collection."""
     m_collection = MetricCollection(
         BinaryAccuracy(),
         BinaryPrecision(),
         BinaryRecall(),
+        prefix=prefix,
+        postfix=postfix,
     )
     if num_vals == 1:
         m_collection.update(torch.randint(0, 2, size=(10,)), torch.randint(0, 2, size=(10,)))