Fix logic mistake in plotting of metric collection (#1783)

* fix logic mistake

* changelog

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Jirka Borovec <6035284+Borda@users.noreply.github.com>

CHANGELOG.md

@@ -189,8 +189,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fixed use of `prefix` and `postfix` in nested `MetricCollection` ([#1773](https://github.com/Lightning-AI/torchmetrics/pull/1773))
 
+
+- Fixed `ax` plotting logging in `MetricCollection ([#1783](https://github.com/Lightning-AI/torchmetrics/pull/1783))
+
+
 - Fixed lookup for punkt sources being downloaded in `RougeScore` [#1789](https://github.com/Lightning-AI/torchmetrics/pull/1789)
 
+
 ## [0.11.4] - 2023-03-10
 
 ### Fixed

src/torchmetrics/collections.py

@@ -593,11 +593,8 @@ def plot(
                 raise ValueError(
                     f"Expected argument `ax` to be a matplotlib axis object, but got {type(ax)} when `together=True`"
                 )
-            if (
-                not together
-                and not isinstance(ax, Sequence)
-                and not all(isinstance(a, _AX_TYPE) for a in ax)
-                and len(ax) != len(self)
+            if not together and not (
+                isinstance(ax, Sequence) and all(isinstance(a, _AX_TYPE) for a in ax) and len(ax) == len(self)
             ):
                 raise ValueError(
                     f"Expected argument `ax` to be a sequence of matplotlib axis objects with the same length as the "

tests/unittests/utilities/test_plot.py

@@ -824,6 +824,14 @@ def test_plot_method_collection(together, num_vals):
         assert all(isinstance(f[0], plt.Figure) for f in fig_ax)
         assert all(isinstance(f[1], matplotlib.axes.Axes) for f in fig_ax)
 
+    # test ax arg
+    fig, ax = plt.subplots(nrows=len(m_collection), ncols=1)
+    m_collection.plot(ax=ax.tolist())
+
+    fig, ax = plt.subplots(nrows=len(m_collection) + 1, ncols=1)
+    with pytest.raises(ValueError, match="Expected argument `ax` to be a sequence of matplotlib axis objects with.*"):
+        m_collection.plot(ax=ax.tolist())
+
 
 @pytest.mark.parametrize(
     ("metric_class", "preds", "target"),