Classification Multilabel Micro AveragePrecision does not form a compute group with comparable metrics

[tsteffek]

🐛 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) while AUROC flattens only in its compute() (see here). Because of that the shapes don't align and the compute group merge will fail.

To Reproduce

Code sample

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.