Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/metrics/metric.py (Python)

Code before:
        >>> metrics(preds, target)
        {'micro_recall': tensor(0.1250), 'weighted_recall': tensor(0.1111)}

    """
    def __init__(self, metrics):
        super().__init__()
        if isinstance(metrics, dict):
            # Check all values are metrics
            for name, metric in metrics.items():
                if not isinstance(metric, Metric):
                    raise ValueError(f'Value {metric} belonging to key {name}'
                                     ' is not an instance of `pl.metrics.Metric`')
                self[name] = metric
        elif isinstance(metrics, (tuple, list)):
            for metric in metrics:
                if not isinstance(metric, Metric):
                    raise ValueError(f'Input {metric} to `MetricCollection` is not a instance'
                                     ' of `pl.metrics.Metric`')
                name = metric.__class__.__name__
                if name in self:
                    raise ValueError(f'Encountered two metrics both named {name}')
                self[name] = metric
        else:
            raise ValueError('Unknown input to MetricCollection.')

    def forward(self, *args, **kwargs):
        """ Iteratively call forward for each metric """
        return {k: m(*args, **kwargs) for k, m in self.items()}

    def update(self, *args, **kwargs):
        """ Iteratively call update for each metric """
        for _, m in self.items():
            m.update(*args, **kwargs)

    def compute(self):
        return {k: m.compute() for k, m in self.items()}

    def reset(self):
        """ Iteratively call reset for each metric """
        for _, m in self.items():
            m.reset()


Review comment (refactor):
let's say I have two metrics like this:
`Metric1` takes inputs `a` and  `b`. `Metric2` takes inputs `c` and `d`. With your implementations both metrics would get `a`, `b`, `c` and `d`, correct? 

I think the cleaner way would be to pass all args and extract kwargs based on the metric's method signature using inspect.

What do you think?

Code after:
                                f'or a list of torch.Tensor, but encountered {current_val}')
        return self

    def persistent(self, mode: bool = False):
        """ Method for post-init to change if metric states should be saved to
            its state_dict
        """
        for key in self._persistent.keys():
            self._persistent[key] = mode

    def state_dict(self, *args, **kwargs):
        # Register metric states to be part of the state_dict
        state_dict = super().state_dict()
        for key in self._defaults.keys():
            if self._persistent[key]:
                current_val = getattr(self, key)
                state_dict.update({key: current_val})
        return state_dict


class MetricCollection(nn.ModuleDict):
    """
    MetricCollection class can be used to chain metrics that have the same
    call pattern into one single class.

    Args:
        metrics: One of the following

            * list or tuple: if metrics are passed in as a list, will use the
              metrics class name as key for output dict. Therefore, two metrics
              of the same class cannot be chained this way.

            * dict: if metrics are passed in as a dict, will use each key in the
              dict as key for output dict. Use this format if you want to chain
              together multiple of the same metric with different parameters.

    Example (input as list):

        >>> from pytorch_lightning.metrics import MetricCollection, Accuracy, Precision, Recall
        >>> target = torch.tensor([0, 2, 0, 2, 0, 1, 0, 2])
        >>> preds = torch.tensor([2, 1, 2, 0, 1, 2, 2, 2])
        >>> metrics = MetricCollection([Accuracy(),
        ...                             Precision(num_classes=3, average='macro'),
        ...                             Recall(num_classes=3, average='macro')])
        >>> metrics(preds, target)
        {'Accuracy': tensor(0.1250), 'Precision': tensor(0.0667), 'Recall': tensor(0.1111)}

    Example (input as dict):

        >>> metrics = MetricCollection({'micro_recall': Recall(num_classes=3, average='micro'),
        ...                             'macro_recall': Recall(num_classes=3, average='macro')})