Repository: Lightning-AI/torchmetrics

Issue:
Classification Multilabel Micro AveragePrecision does not form a compute group with comparable metrics
## 🐛 Bug

While micro and macro `AUROC` play well with each other and macro `AveragePrecision`, micro `AveragePrecision` will not be merged into the same compute group.

This is due to AveragePrecision flattening its predictions and targets in the `update()` call ([see here](https://github.com/PyTorchLightning/metrics/blob/ea97d033ea3c7d77a927907c4d6232f1ada6469a/src/torchmetrics/functional/classification/average_precision.py#L46-L51)) while AUROC flattens only in its `compute()` ([see here](https://github.com/PyTorchLightning/metrics/blob/ea97d033ea3c7d77a927907c4d6232f1ada6469a/src/torchmetrics/functional/classification/auroc.py#L120-L121)). Because of that the shapes don't align and the compute group merge will fail.

### To Reproduce

#### Code sample

```python
import torch
from torchmetrics import MetricCollection, AUROC, AveragePrecision

m = MetricCollection([MetricCollection(AUROC(average='micro', num_classes=3), AveragePrecision(average='micro', num_classes=3), postfix='_micro'),
                     MetricCollection(AUROC(average='macro', num_classes=3), AveragePrecision(average='macro', num_classes=3), postfix='_macro')])
# Multi-label inputs
ml_preds  = torch.tensor([[0.2, 0.8, 0.9], [0.5, 0.6, 0.1], [0.3, 0.1, 0.1]])
ml_target = torch.tensor([[0, 1, 1], [1, 0, 0], [0, 0, 0]])

m._groups
# Out: 
# {0: ['AUROC_micro'],
#  1: ['AveragePrecision_micro'],
#  2: ['AUROC_macro'],
#  3: ['AveragePrecision_macro']}
m.update(ml_preds, ml_target)
m._groups
# Out: 
# {0: ['AUROC_micro', 'AUROC_macro', 'AveragePrecision_macro'],
#  1: ['AveragePrecision_micro']} - maybe `AveragePrecision_micro` has body odor?
```

### Expected behavior

Micro `AveragePrecision` shouldn't flatten during update but during compute, which would allow it to have its state shared with e.g. AUROC and itself.

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source): `0.9.1`, `pip`
- Python & PyTorch Version (e.g., 1.0):  `3.8.12` & `1.11.0`
- Any other relevant information such as OS (e.g., Linux): `FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel`

### Additional context

This is especially hurting for AveragePrecision and the likes, since they store all predictions and targets in their state. 



Discussion:

Resolving patch:
diff --git a/src/torchmetrics/functional/classification/average_precision.py b/src/torchmetrics/functional/classification/average_precision.py
--- a/src/torchmetrics/functional/classification/average_precision.py
+++ b/src/torchmetrics/functional/classification/average_precision.py
@@ -43,14 +43,8 @@ def _average_precision_update(
         average: reduction method for multi-class or multi-label problems
     """
     preds, target, num_classes, pos_label = _precision_recall_curve_update(preds, target, num_classes, pos_label)
-    if average == "micro":
-        if preds.ndim == target.ndim:
-            # Considering each element of the label indicator matrix as a label
-            preds = preds.flatten()
-            target = target.flatten()
-            num_classes = 1
-        else:
-            raise ValueError("Cannot use `micro` average with multi-class input")
+    if average == "micro" and preds.ndim != target.ndim:
+        raise ValueError("Cannot use `micro` average with multi-class input")
 
     return preds, target, num_classes, pos_label
 
@@ -97,6 +91,11 @@ def _average_precision_compute(
     """
 
     # todo: `sample_weights` is unused
+    if average == "micro" and preds.ndim == target.ndim:
+        preds = preds.flatten()
+        target = target.flatten()
+        num_classes = 1
+
     precision, recall, _ = _precision_recall_curve_compute(preds, target, num_classes, pos_label)
     if average == "weighted":
         if preds.ndim == target.ndim and target.ndim > 1:


Test patch:
diff --git a/test/unittests/bases/test_collections.py b/test/unittests/bases/test_collections.py
--- a/test/unittests/bases/test_collections.py
+++ b/test/unittests/bases/test_collections.py
@@ -19,7 +19,9 @@
 import torch
 
 from torchmetrics import (
+    AUROC,
     Accuracy,
+    AveragePrecision,
     CohenKappa,
     ConfusionMatrix,
     F1Score,
@@ -29,6 +31,7 @@
     Precision,
     Recall,
 )
+from torchmetrics.utilities.checks import _allclose_recursive
 from unittests.helpers import seed_all
 from unittests.helpers.testers import DummyMetricDiff, DummyMetricSum
 
@@ -267,6 +270,8 @@ def test_collection_filtering():
     """Test that collections works with the kwargs argument."""
 
     class DummyMetric(Metric):
+        full_state_update = True
+
         def __init__(self):
             super().__init__()
 
@@ -277,6 +282,8 @@ def compute(self):
             return
 
     class MyAccuracy(Metric):
+        full_state_update = True
+
         def __init__(self):
             super().__init__()
 
@@ -292,21 +299,30 @@ def compute(self):
     mc2(torch.tensor([0, 1]), torch.tensor([0, 1]), kwarg="kwarg", kwarg2="kwarg2")
 
 
+# function for generating
+_mc_preds = torch.randn(10, 3).softmax(dim=-1)
+_mc_target = torch.randint(3, (10,))
+_ml_preds = torch.rand(10, 3)
+_ml_target = torch.randint(2, (10, 3))
+
+
 @pytest.mark.parametrize(
-    "metrics, expected",
+    "metrics, expected, preds, target",
     [
         # single metric forms its own compute group
-        (Accuracy(3), {0: ["Accuracy"]}),
+        (Accuracy(3), {0: ["Accuracy"]}, _mc_preds, _mc_target),
         # two metrics of same class forms a compute group
-        ({"acc0": Accuracy(3), "acc1": Accuracy(3)}, {0: ["acc0", "acc1"]}),
+        ({"acc0": Accuracy(3), "acc1": Accuracy(3)}, {0: ["acc0", "acc1"]}, _mc_preds, _mc_target),
         # two metrics from registry froms a compute group
-        ([Precision(3), Recall(3)], {0: ["Precision", "Recall"]}),
+        ([Precision(3), Recall(3)], {0: ["Precision", "Recall"]}, _mc_preds, _mc_target),
         # two metrics from different classes gives two compute groups
-        ([ConfusionMatrix(3), Recall(3)], {0: ["ConfusionMatrix"], 1: ["Recall"]}),
+        ([ConfusionMatrix(3), Recall(3)], {0: ["ConfusionMatrix"], 1: ["Recall"]}, _mc_preds, _mc_target),
         # multi group multi metric
         (
             [ConfusionMatrix(3), CohenKappa(3), Recall(3), Precision(3)],
             {0: ["ConfusionMatrix", "CohenKappa"], 1: ["Recall", "Precision"]},
+            _mc_preds,
+            _mc_target,
         ),
         # Complex example
         (
@@ -319,6 +335,33 @@ def compute(self):
                 "confmat": ConfusionMatrix(3),
             },
             {0: ["acc", "acc2", "f1", "recall"], 1: ["acc3"], 2: ["confmat"]},
+            _mc_preds,
+            _mc_target,
+        ),
+        # With list states
+        (
+            [AUROC(average="macro", num_classes=3), AveragePrecision(average="macro", num_classes=3)],
+            {0: ["AUROC", "AveragePrecision"]},
+            _mc_preds,
+            _mc_target,
+        ),
+        # Nested collections
+        (
+            [
+                MetricCollection(
+                    AUROC(average="micro", num_classes=3),
+                    AveragePrecision(average="micro", num_classes=3),
+                    postfix="_micro",
+                ),
+                MetricCollection(
+                    AUROC(average="macro", num_classes=3),
+                    AveragePrecision(average="macro", num_classes=3),
+                    postfix="_macro",
+                ),
+            ],
+            {0: ["AUROC_micro", "AveragePrecision_micro", "AUROC_macro", "AveragePrecision_macro"]},
+            _ml_preds,
+            _ml_target,
         ),
     ],
 )
@@ -332,8 +375,10 @@ class TestComputeGroups:
             ["prefix_", "_postfix"],
         ],
     )
-    def test_check_compute_groups_correctness(self, metrics, expected, prefix, postfix):
+    def test_check_compute_groups_correctness(self, metrics, expected, preds, target, prefix, postfix):
         """Check that compute groups are formed after initialization and that metrics are correctly computed."""
+        if isinstance(metrics, MetricCollection):
+            prefix, postfix = None, None  # disable for nested collections
         m = MetricCollection(deepcopy(metrics), prefix=prefix, postfix=postfix, compute_groups=True)
         # Construct without for comparison
         m2 = MetricCollection(deepcopy(metrics), prefix=prefix, postfix=postfix, compute_groups=False)
@@ -342,8 +387,6 @@ def test_check_compute_groups_correctness(self, metrics, expected, prefix, postf
         assert m2.compute_groups == {}
 
         for _ in range(2):  # repeat to emulate effect of multiple epochs
-            preds = torch.randn(10, 3).softmax(dim=-1)
-            target = torch.randint(3, (10,))
             m.update(preds, target)
             m2.update(preds, target)
 
@@ -353,8 +396,6 @@ def test_check_compute_groups_correctness(self, metrics, expected, prefix, postf
             assert m.compute_groups == expected
             assert m2.compute_groups == {}
 
-            preds = torch.randn(10, 3).softmax(dim=-1)
-            target = torch.randint(3, (10,))
             # compute groups should kick in here
             m.update(preds, target)
             m2.update(preds, target)
@@ -372,7 +413,7 @@ def test_check_compute_groups_correctness(self, metrics, expected, prefix, postf
             m2.reset()
 
     @pytest.mark.parametrize("method", ["items", "values", "keys"])
-    def test_check_compute_groups_items_and_values(self, metrics, expected, method):
+    def test_check_compute_groups_items_and_values(self, metrics, expected, preds, target, method):
         """Check that whenever user call a methods that give access to the indivitual metric that state are copied
         instead of just passed by reference."""
         m = MetricCollection(deepcopy(metrics), compute_groups=True)
@@ -380,14 +421,12 @@ def test_check_compute_groups_items_and_values(self, metrics, expected, method):
 
         for _ in range(2):  # repeat to emulate effect of multiple epochs
             for _ in range(2):  # repeat to emulate effect of multiple batches
-                preds = torch.randn(10, 3).softmax(dim=-1)
-                target = torch.randint(3, (10,))
                 m.update(preds, target)
                 m2.update(preds, target)
 
             def _compare(m1, m2):
                 for state in m1._defaults:
-                    assert torch.allclose(getattr(m1, state), getattr(m2, state))
+                    assert _allclose_recursive(getattr(m1, state), getattr(m2, state))
                 # if states are still by reference the reset will make following metrics fail
                 m1.reset()
                 m2.reset()
diff --git a/test/unittests/bases/test_composition.py b/test/unittests/bases/test_composition.py
--- a/test/unittests/bases/test_composition.py
+++ b/test/unittests/bases/test_composition.py
@@ -22,6 +22,8 @@
 
 
 class DummyMetric(Metric):
+    full_state_update = True
+
     def __init__(self, val_to_return):
         super().__init__()
         self.add_state("_num_updates", tensor(0), dist_reduce_fx="sum")