We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Can MultitaskWrapper.clone() have a postfix or prefix arg?
## 🚀 Feature
A `prefix` and `postfix` parameter to pass into the `clone()` method of `MultitaskWrapper`

### Motivation

I have fabricated a `MultitaskWrapper` to measure multiple task metrics for trainning, and wanted to clone it for validation and test.
```
self.metrics = MultitaskWrapper({
            'Img': MetricCollection({
                'PSNR': PeakSignalNoiseRatio(data_range=1.0),
                'SSIM': StructuralSimilarityIndexMeasure(data_range=1.0),
            }),
            'DoP': MetricCollection({`some_metrics`: SomeMetrics()}),
            'AoP': MetricCollection({`other_metrics`: OtherMetrics()})
        })
```
What I want to do is:
`self.val_metrics =  self.metrics.clone(postfix='_val')` so that we can obtain
```
self.val_metrics = MultitaskWrapper({
            'Img_val': MetricCollection({
                'PSNR': PeakSignalNoiseRatio(data_range=1.0),
                'SSIM': StructuralSimilarityIndexMeasure(data_range=1.0),
            }),
            'DoP_val': MetricCollection({`some_metrics`: SomeMetrics()}),
            'AoP_val': MetricCollection({`other_metrics`: OtherMetrics()})
        })
```
However, there is no such a postfix arg for `clone()`. So I can only re-define a `MultitaskWrapper` for `self.val_metrics` instead of simply cloning form `self.metrics`.

Is there any way to handle it?
Thx.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/torchmetrics/wrappers/multitask.py`
Content:
```
1 # Copyright The Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # this is just a bypass for this module name collision with built-in one
15 from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
16 
17 from torch import Tensor, nn
18 
19 from torchmetrics.collections import MetricCollection
20 from torchmetrics.metric import Metric
21 from torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE
22 from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE
23 from torchmetrics.wrappers.abstract import WrapperMetric
24 
25 if not _MATPLOTLIB_AVAILABLE:
26     __doctest_skip__ = ["MultitaskWrapper.plot"]
27 
28 
29 class MultitaskWrapper(WrapperMetric):
30     """Wrapper class for computing different metrics on different tasks in the context of multitask learning.
31 
32     In multitask learning the different tasks requires different metrics to be evaluated. This wrapper allows
33     for easy evaluation in such cases by supporting multiple predictions and targets through a dictionary.
34     Note that only metrics where the signature of `update` follows the standard `preds, target` is supported.
35 
36     Args:
37         task_metrics:
38             Dictionary associating each task to a Metric or a MetricCollection. The keys of the dictionary represent the
39             names of the tasks, and the values represent the metrics to use for each task.
40 
41     Raises:
42         TypeError:
43             If argument `task_metrics` is not an dictionary
44         TypeError:
45             If not all values in the `task_metrics` dictionary is instances of `Metric` or `MetricCollection`
46 
47     Example (with a single metric per class):
48          >>> import torch
49          >>> from torchmetrics.wrappers import MultitaskWrapper
50          >>> from torchmetrics.regression import MeanSquaredError
51          >>> from torchmetrics.classification import BinaryAccuracy
52          >>>
53          >>> classification_target = torch.tensor([0, 1, 0])
54          >>> regression_target = torch.tensor([2.5, 5.0, 4.0])
55          >>> targets = {"Classification": classification_target, "Regression": regression_target}
56          >>>
57          >>> classification_preds = torch.tensor([0, 0, 1])
58          >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
59          >>> preds = {"Classification": classification_preds, "Regression": regression_preds}
60          >>>
61          >>> metrics = MultitaskWrapper({
62          ...     "Classification": BinaryAccuracy(),
63          ...     "Regression": MeanSquaredError()
64          ... })
65          >>> metrics.update(preds, targets)
66          >>> metrics.compute()
67          {'Classification': tensor(0.3333), 'Regression': tensor(0.8333)}
68 
69     Example (with several metrics per task):
70          >>> import torch
71          >>> from torchmetrics import MetricCollection
72          >>> from torchmetrics.wrappers import MultitaskWrapper
73          >>> from torchmetrics.regression import MeanSquaredError, MeanAbsoluteError
74          >>> from torchmetrics.classification import BinaryAccuracy, BinaryF1Score
75          >>>
76          >>> classification_target = torch.tensor([0, 1, 0])
77          >>> regression_target = torch.tensor([2.5, 5.0, 4.0])
78          >>> targets = {"Classification": classification_target, "Regression": regression_target}
79          >>>
80          >>> classification_preds = torch.tensor([0, 0, 1])
81          >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
82          >>> preds = {"Classification": classification_preds, "Regression": regression_preds}
83          >>>
84          >>> metrics = MultitaskWrapper({
85          ...     "Classification": MetricCollection(BinaryAccuracy(), BinaryF1Score()),
86          ...     "Regression": MetricCollection(MeanSquaredError(), MeanAbsoluteError())
87          ... })
88          >>> metrics.update(preds, targets)
89          >>> metrics.compute()
90          {'Classification': {'BinaryAccuracy': tensor(0.3333), 'BinaryF1Score': tensor(0.)},
91           'Regression': {'MeanSquaredError': tensor(0.8333), 'MeanAbsoluteError': tensor(0.6667)}}
92 
93     """
94 
95     is_differentiable = False
96 
97     def __init__(
98         self,
99         task_metrics: Dict[str, Union[Metric, MetricCollection]],
100     ) -> None:
101         self._check_task_metrics_type(task_metrics)
102         super().__init__()
103         self.task_metrics = nn.ModuleDict(task_metrics)
104 
105     def items(self) -> Iterable[Tuple[str, nn.Module]]:
106         """Iterate over task and task metrics."""
107         return self.task_metrics.items()
108 
109     def keys(self) -> Iterable[str]:
110         """Iterate over task names."""
111         return self.task_metrics.keys()
112 
113     def values(self) -> Iterable[nn.Module]:
114         """Iterate over task metrics."""
115         return self.task_metrics.values()
116 
117     @staticmethod
118     def _check_task_metrics_type(task_metrics: Dict[str, Union[Metric, MetricCollection]]) -> None:
119         if not isinstance(task_metrics, dict):
120             raise TypeError(f"Expected argument `task_metrics` to be a dict. Found task_metrics = {task_metrics}")
121 
122         for metric in task_metrics.values():
123             if not (isinstance(metric, (Metric, MetricCollection))):
124                 raise TypeError(
125                     "Expected each task's metric to be a Metric or a MetricCollection. "
126                     f"Found a metric of type {type(metric)}"
127                 )
128 
129     def update(self, task_preds: Dict[str, Tensor], task_targets: Dict[str, Tensor]) -> None:
130         """Update each task's metric with its corresponding pred and target.
131 
132         Args:
133             task_preds: Dictionary associating each task to a Tensor of pred.
134             task_targets: Dictionary associating each task to a Tensor of target.
135 
136         """
137         if not self.task_metrics.keys() == task_preds.keys() == task_targets.keys():
138             raise ValueError(
139                 "Expected arguments `task_preds` and `task_targets` to have the same keys as the wrapped `task_metrics`"
140                 f". Found task_preds.keys() = {task_preds.keys()}, task_targets.keys() = {task_targets.keys()} "
141                 f"and self.task_metrics.keys() = {self.task_metrics.keys()}"
142             )
143 
144         for task_name, metric in self.task_metrics.items():
145             pred = task_preds[task_name]
146             target = task_targets[task_name]
147             metric.update(pred, target)
148 
149     def compute(self) -> Dict[str, Any]:
150         """Compute metrics for all tasks."""
151         return {task_name: metric.compute() for task_name, metric in self.task_metrics.items()}
152 
153     def forward(self, task_preds: Dict[str, Tensor], task_targets: Dict[str, Tensor]) -> Dict[str, Any]:
154         """Call underlying forward methods for all tasks and return the result as a dictionary."""
155         # This method is overridden because we do not need the complex version defined in Metric, that relies on the
156         # value of full_state_update, and that also accumulates the results. Here, all computations are handled by the
157         # underlying metrics, which all have their own value of full_state_update, and which all accumulate the results
158         # by themselves.
159         return {
160             task_name: metric(task_preds[task_name], task_targets[task_name])
161             for task_name, metric in self.task_metrics.items()
162         }
163 
164     def reset(self) -> None:
165         """Reset all underlying metrics."""
166         for metric in self.task_metrics.values():
167             metric.reset()
168         super().reset()
169 
170     def plot(
171         self, val: Optional[Union[Dict, Sequence[Dict]]] = None, axes: Optional[Sequence[_AX_TYPE]] = None
172     ) -> Sequence[_PLOT_OUT_TYPE]:
173         """Plot a single or multiple values from the metric.
174 
175         All tasks' results are plotted on individual axes.
176 
177         Args:
178             val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.
179                 If no value is provided, will automatically call `metric.compute` and plot that result.
180             axes: Sequence of matplotlib axis objects. If provided, will add the plots to the provided axis objects.
181                 If not provided, will create them.
182 
183         Returns:
184             Sequence of tuples with Figure and Axes object for each task.
185 
186         .. plot::
187             :scale: 75
188 
189             >>> # Example plotting a single value
190             >>> import torch
191             >>> from torchmetrics.wrappers import MultitaskWrapper
192             >>> from torchmetrics.regression import MeanSquaredError
193             >>> from torchmetrics.classification import BinaryAccuracy
194             >>>
195             >>> classification_target = torch.tensor([0, 1, 0])
196             >>> regression_target = torch.tensor([2.5, 5.0, 4.0])
197             >>> targets = {"Classification": classification_target, "Regression": regression_target}
198             >>>
199             >>> classification_preds = torch.tensor([0, 0, 1])
200             >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
201             >>> preds = {"Classification": classification_preds, "Regression": regression_preds}
202             >>>
203             >>> metrics = MultitaskWrapper({
204             ...     "Classification": BinaryAccuracy(),
205             ...     "Regression": MeanSquaredError()
206             ... })
207             >>> metrics.update(preds, targets)
208             >>> value = metrics.compute()
209             >>> fig_, ax_ = metrics.plot(value)
210 
211         .. plot::
212             :scale: 75
213 
214             >>> # Example plotting multiple values
215             >>> import torch
216             >>> from torchmetrics.wrappers import MultitaskWrapper
217             >>> from torchmetrics.regression import MeanSquaredError
218             >>> from torchmetrics.classification import BinaryAccuracy
219             >>>
220             >>> classification_target = torch.tensor([0, 1, 0])
221             >>> regression_target = torch.tensor([2.5, 5.0, 4.0])
222             >>> targets = {"Classification": classification_target, "Regression": regression_target}
223             >>>
224             >>> classification_preds = torch.tensor([0, 0, 1])
225             >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])
226             >>> preds = {"Classification": classification_preds, "Regression": regression_preds}
227             >>>
228             >>> metrics = MultitaskWrapper({
229             ...     "Classification": BinaryAccuracy(),
230             ...     "Regression": MeanSquaredError()
231             ... })
232             >>> values = []
233             >>> for _ in range(10):
234             ...     values.append(metrics(preds, targets))
235             >>> fig_, ax_ = metrics.plot(values)
236 
237         """
238         if axes is not None:
239             if not isinstance(axes, Sequence):
240                 raise TypeError(f"Expected argument `axes` to be a Sequence. Found type(axes) = {type(axes)}")
241 
242             if not all(isinstance(ax, _AX_TYPE) for ax in axes):
243                 raise TypeError("Expected each ax in argument `axes` to be a matplotlib axis object")
244 
245             if len(axes) != len(self.task_metrics):
246                 raise ValueError(
247                     "Expected argument `axes` to be a Sequence of the same length as the number of tasks."
248                     f"Found len(axes) = {len(axes)} and {len(self.task_metrics)} tasks"
249                 )
250 
251         val = val if val is not None else self.compute()
252         fig_axs = []
253         for i, (task_name, task_metric) in enumerate(self.task_metrics.items()):
254             ax = axes[i] if axes is not None else None
255             if isinstance(val, Dict):
256                 f, a = task_metric.plot(val[task_name], ax=ax)
257             elif isinstance(val, Sequence):
258                 f, a = task_metric.plot([v[task_name] for v in val], ax=ax)
259             else:
260                 raise TypeError(
261                     "Expected argument `val` to be None or of type Dict or Sequence[Dict]. "
262                     f"Found type(val)= {type(val)}"
263                 )
264             fig_axs.append((f, a))
265         return fig_axs
266
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/torchmetrics/wrappers/multitask.py b/src/torchmetrics/wrappers/multitask.py
--- a/src/torchmetrics/wrappers/multitask.py
+++ b/src/torchmetrics/wrappers/multitask.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # this is just a bypass for this module name collision with built-in one
+from copy import deepcopy
 from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union
 
 from torch import Tensor, nn
@@ -167,6 +168,33 @@
             metric.reset()
         super().reset()
 
+    @staticmethod
+    def _check_arg(arg: Optional[str], name: str) -> Optional[str]:
+        if arg is None or isinstance(arg, str):
+            return arg
+        raise ValueError(f"Expected input `{name}` to be a string, but got {type(arg)}")
+
+    def clone(self, prefix: Optional[str] = None, postfix: Optional[str] = None) -> "MultitaskWrapper":
+        """Make a copy of the metric.
+
+        Args:
+            prefix: a string to append in front of the metric keys
+            postfix: a string to append after the keys of the output dict.
+
+        """
+        multitask_copy = deepcopy(self)
+        if prefix is not None:
+            prefix = self._check_arg(prefix, "prefix")
+            multitask_copy.task_metrics = nn.ModuleDict(
+                {prefix + key: value for key, value in multitask_copy.task_metrics.items()}
+            )
+        if postfix is not None:
+            postfix = self._check_arg(postfix, "postfix")
+            multitask_copy.task_metrics = nn.ModuleDict(
+                {key + postfix: value for key, value in multitask_copy.task_metrics.items()}
+            )
+        return multitask_copy
+
     def plot(
         self, val: Optional[Union[Dict, Sequence[Dict]]] = None, axes: Optional[Sequence[_AX_TYPE]] = None
     ) -> Sequence[_PLOT_OUT_TYPE]:

{"golden_diff": "diff --git a/src/torchmetrics/wrappers/multitask.py b/src/torchmetrics/wrappers/multitask.py\n--- a/src/torchmetrics/wrappers/multitask.py\n+++ b/src/torchmetrics/wrappers/multitask.py\n@@ -12,6 +12,7 @@\n # See the License for the specific language governing permissions and\n # limitations under the License.\n # this is just a bypass for this module name collision with built-in one\n+from copy import deepcopy\n from typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union\n \n from torch import Tensor, nn\n@@ -167,6 +168,33 @@\n             metric.reset()\n         super().reset()\n \n+    @staticmethod\n+    def _check_arg(arg: Optional[str], name: str) -> Optional[str]:\n+        if arg is None or isinstance(arg, str):\n+            return arg\n+        raise ValueError(f\"Expected input `{name}` to be a string, but got {type(arg)}\")\n+\n+    def clone(self, prefix: Optional[str] = None, postfix: Optional[str] = None) -> \"MultitaskWrapper\":\n+        \"\"\"Make a copy of the metric.\n+\n+        Args:\n+            prefix: a string to append in front of the metric keys\n+            postfix: a string to append after the keys of the output dict.\n+\n+        \"\"\"\n+        multitask_copy = deepcopy(self)\n+        if prefix is not None:\n+            prefix = self._check_arg(prefix, \"prefix\")\n+            multitask_copy.task_metrics = nn.ModuleDict(\n+                {prefix + key: value for key, value in multitask_copy.task_metrics.items()}\n+            )\n+        if postfix is not None:\n+            postfix = self._check_arg(postfix, \"postfix\")\n+            multitask_copy.task_metrics = nn.ModuleDict(\n+                {key + postfix: value for key, value in multitask_copy.task_metrics.items()}\n+            )\n+        return multitask_copy\n+\n     def plot(\n         self, val: Optional[Union[Dict, Sequence[Dict]]] = None, axes: Optional[Sequence[_AX_TYPE]] = None\n     ) -> Sequence[_PLOT_OUT_TYPE]:\n", "issue": "Can MultitaskWrapper.clone() have a postfix or prefix arg?\n## \ud83d\ude80 Feature\r\nA `prefix` and `postfix` parameter to pass into the `clone()` method of `MultitaskWrapper`\r\n\r\n### Motivation\r\n\r\nI have fabricated a `MultitaskWrapper` to measure multiple task metrics for trainning, and wanted to clone it for validation and test.\r\n```\r\nself.metrics = MultitaskWrapper({\r\n            'Img': MetricCollection({\r\n                'PSNR': PeakSignalNoiseRatio(data_range=1.0),\r\n                'SSIM': StructuralSimilarityIndexMeasure(data_range=1.0),\r\n            }),\r\n            'DoP': MetricCollection({`some_metrics`: SomeMetrics()}),\r\n            'AoP': MetricCollection({`other_metrics`: OtherMetrics()})\r\n        })\r\n```\r\nWhat I want to do is:\r\n`self.val_metrics =  self.metrics.clone(postfix='_val')` so that we can obtain\r\n```\r\nself.val_metrics = MultitaskWrapper({\r\n            'Img_val': MetricCollection({\r\n                'PSNR': PeakSignalNoiseRatio(data_range=1.0),\r\n                'SSIM': StructuralSimilarityIndexMeasure(data_range=1.0),\r\n            }),\r\n            'DoP_val': MetricCollection({`some_metrics`: SomeMetrics()}),\r\n            'AoP_val': MetricCollection({`other_metrics`: OtherMetrics()})\r\n        })\r\n```\r\nHowever, there is no such a postfix arg for `clone()`. So I can only re-define a `MultitaskWrapper` for `self.val_metrics` instead of simply cloning form `self.metrics`.\r\n\r\nIs there any way to handle it?\r\nThx.\r\n\n", "before_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# this is just a bypass for this module name collision with built-in one\nfrom typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union\n\nfrom torch import Tensor, nn\n\nfrom torchmetrics.collections import MetricCollection\nfrom torchmetrics.metric import Metric\nfrom torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE\nfrom torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE\nfrom torchmetrics.wrappers.abstract import WrapperMetric\n\nif not _MATPLOTLIB_AVAILABLE:\n    __doctest_skip__ = [\"MultitaskWrapper.plot\"]\n\n\nclass MultitaskWrapper(WrapperMetric):\n    \"\"\"Wrapper class for computing different metrics on different tasks in the context of multitask learning.\n\n    In multitask learning the different tasks requires different metrics to be evaluated. This wrapper allows\n    for easy evaluation in such cases by supporting multiple predictions and targets through a dictionary.\n    Note that only metrics where the signature of `update` follows the standard `preds, target` is supported.\n\n    Args:\n        task_metrics:\n            Dictionary associating each task to a Metric or a MetricCollection. The keys of the dictionary represent the\n            names of the tasks, and the values represent the metrics to use for each task.\n\n    Raises:\n        TypeError:\n            If argument `task_metrics` is not an dictionary\n        TypeError:\n            If not all values in the `task_metrics` dictionary is instances of `Metric` or `MetricCollection`\n\n    Example (with a single metric per class):\n         >>> import torch\n         >>> from torchmetrics.wrappers import MultitaskWrapper\n         >>> from torchmetrics.regression import MeanSquaredError\n         >>> from torchmetrics.classification import BinaryAccuracy\n         >>>\n         >>> classification_target = torch.tensor([0, 1, 0])\n         >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n         >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n         >>>\n         >>> classification_preds = torch.tensor([0, 0, 1])\n         >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n         >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n         >>>\n         >>> metrics = MultitaskWrapper({\n         ...     \"Classification\": BinaryAccuracy(),\n         ...     \"Regression\": MeanSquaredError()\n         ... })\n         >>> metrics.update(preds, targets)\n         >>> metrics.compute()\n         {'Classification': tensor(0.3333), 'Regression': tensor(0.8333)}\n\n    Example (with several metrics per task):\n         >>> import torch\n         >>> from torchmetrics import MetricCollection\n         >>> from torchmetrics.wrappers import MultitaskWrapper\n         >>> from torchmetrics.regression import MeanSquaredError, MeanAbsoluteError\n         >>> from torchmetrics.classification import BinaryAccuracy, BinaryF1Score\n         >>>\n         >>> classification_target = torch.tensor([0, 1, 0])\n         >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n         >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n         >>>\n         >>> classification_preds = torch.tensor([0, 0, 1])\n         >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n         >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n         >>>\n         >>> metrics = MultitaskWrapper({\n         ...     \"Classification\": MetricCollection(BinaryAccuracy(), BinaryF1Score()),\n         ...     \"Regression\": MetricCollection(MeanSquaredError(), MeanAbsoluteError())\n         ... })\n         >>> metrics.update(preds, targets)\n         >>> metrics.compute()\n         {'Classification': {'BinaryAccuracy': tensor(0.3333), 'BinaryF1Score': tensor(0.)},\n          'Regression': {'MeanSquaredError': tensor(0.8333), 'MeanAbsoluteError': tensor(0.6667)}}\n\n    \"\"\"\n\n    is_differentiable = False\n\n    def __init__(\n        self,\n        task_metrics: Dict[str, Union[Metric, MetricCollection]],\n    ) -> None:\n        self._check_task_metrics_type(task_metrics)\n        super().__init__()\n        self.task_metrics = nn.ModuleDict(task_metrics)\n\n    def items(self) -> Iterable[Tuple[str, nn.Module]]:\n        \"\"\"Iterate over task and task metrics.\"\"\"\n        return self.task_metrics.items()\n\n    def keys(self) -> Iterable[str]:\n        \"\"\"Iterate over task names.\"\"\"\n        return self.task_metrics.keys()\n\n    def values(self) -> Iterable[nn.Module]:\n        \"\"\"Iterate over task metrics.\"\"\"\n        return self.task_metrics.values()\n\n    @staticmethod\n    def _check_task_metrics_type(task_metrics: Dict[str, Union[Metric, MetricCollection]]) -> None:\n        if not isinstance(task_metrics, dict):\n            raise TypeError(f\"Expected argument `task_metrics` to be a dict. Found task_metrics = {task_metrics}\")\n\n        for metric in task_metrics.values():\n            if not (isinstance(metric, (Metric, MetricCollection))):\n                raise TypeError(\n                    \"Expected each task's metric to be a Metric or a MetricCollection. \"\n                    f\"Found a metric of type {type(metric)}\"\n                )\n\n    def update(self, task_preds: Dict[str, Tensor], task_targets: Dict[str, Tensor]) -> None:\n        \"\"\"Update each task's metric with its corresponding pred and target.\n\n        Args:\n            task_preds: Dictionary associating each task to a Tensor of pred.\n            task_targets: Dictionary associating each task to a Tensor of target.\n\n        \"\"\"\n        if not self.task_metrics.keys() == task_preds.keys() == task_targets.keys():\n            raise ValueError(\n                \"Expected arguments `task_preds` and `task_targets` to have the same keys as the wrapped `task_metrics`\"\n                f\". Found task_preds.keys() = {task_preds.keys()}, task_targets.keys() = {task_targets.keys()} \"\n                f\"and self.task_metrics.keys() = {self.task_metrics.keys()}\"\n            )\n\n        for task_name, metric in self.task_metrics.items():\n            pred = task_preds[task_name]\n            target = task_targets[task_name]\n            metric.update(pred, target)\n\n    def compute(self) -> Dict[str, Any]:\n        \"\"\"Compute metrics for all tasks.\"\"\"\n        return {task_name: metric.compute() for task_name, metric in self.task_metrics.items()}\n\n    def forward(self, task_preds: Dict[str, Tensor], task_targets: Dict[str, Tensor]) -> Dict[str, Any]:\n        \"\"\"Call underlying forward methods for all tasks and return the result as a dictionary.\"\"\"\n        # This method is overridden because we do not need the complex version defined in Metric, that relies on the\n        # value of full_state_update, and that also accumulates the results. Here, all computations are handled by the\n        # underlying metrics, which all have their own value of full_state_update, and which all accumulate the results\n        # by themselves.\n        return {\n            task_name: metric(task_preds[task_name], task_targets[task_name])\n            for task_name, metric in self.task_metrics.items()\n        }\n\n    def reset(self) -> None:\n        \"\"\"Reset all underlying metrics.\"\"\"\n        for metric in self.task_metrics.values():\n            metric.reset()\n        super().reset()\n\n    def plot(\n        self, val: Optional[Union[Dict, Sequence[Dict]]] = None, axes: Optional[Sequence[_AX_TYPE]] = None\n    ) -> Sequence[_PLOT_OUT_TYPE]:\n        \"\"\"Plot a single or multiple values from the metric.\n\n        All tasks' results are plotted on individual axes.\n\n        Args:\n            val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.\n                If no value is provided, will automatically call `metric.compute` and plot that result.\n            axes: Sequence of matplotlib axis objects. If provided, will add the plots to the provided axis objects.\n                If not provided, will create them.\n\n        Returns:\n            Sequence of tuples with Figure and Axes object for each task.\n\n        .. plot::\n            :scale: 75\n\n            >>> # Example plotting a single value\n            >>> import torch\n            >>> from torchmetrics.wrappers import MultitaskWrapper\n            >>> from torchmetrics.regression import MeanSquaredError\n            >>> from torchmetrics.classification import BinaryAccuracy\n            >>>\n            >>> classification_target = torch.tensor([0, 1, 0])\n            >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n            >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n            >>>\n            >>> classification_preds = torch.tensor([0, 0, 1])\n            >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n            >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n            >>>\n            >>> metrics = MultitaskWrapper({\n            ...     \"Classification\": BinaryAccuracy(),\n            ...     \"Regression\": MeanSquaredError()\n            ... })\n            >>> metrics.update(preds, targets)\n            >>> value = metrics.compute()\n            >>> fig_, ax_ = metrics.plot(value)\n\n        .. plot::\n            :scale: 75\n\n            >>> # Example plotting multiple values\n            >>> import torch\n            >>> from torchmetrics.wrappers import MultitaskWrapper\n            >>> from torchmetrics.regression import MeanSquaredError\n            >>> from torchmetrics.classification import BinaryAccuracy\n            >>>\n            >>> classification_target = torch.tensor([0, 1, 0])\n            >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n            >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n            >>>\n            >>> classification_preds = torch.tensor([0, 0, 1])\n            >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n            >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n            >>>\n            >>> metrics = MultitaskWrapper({\n            ...     \"Classification\": BinaryAccuracy(),\n            ...     \"Regression\": MeanSquaredError()\n            ... })\n            >>> values = []\n            >>> for _ in range(10):\n            ...     values.append(metrics(preds, targets))\n            >>> fig_, ax_ = metrics.plot(values)\n\n        \"\"\"\n        if axes is not None:\n            if not isinstance(axes, Sequence):\n                raise TypeError(f\"Expected argument `axes` to be a Sequence. Found type(axes) = {type(axes)}\")\n\n            if not all(isinstance(ax, _AX_TYPE) for ax in axes):\n                raise TypeError(\"Expected each ax in argument `axes` to be a matplotlib axis object\")\n\n            if len(axes) != len(self.task_metrics):\n                raise ValueError(\n                    \"Expected argument `axes` to be a Sequence of the same length as the number of tasks.\"\n                    f\"Found len(axes) = {len(axes)} and {len(self.task_metrics)} tasks\"\n                )\n\n        val = val if val is not None else self.compute()\n        fig_axs = []\n        for i, (task_name, task_metric) in enumerate(self.task_metrics.items()):\n            ax = axes[i] if axes is not None else None\n            if isinstance(val, Dict):\n                f, a = task_metric.plot(val[task_name], ax=ax)\n            elif isinstance(val, Sequence):\n                f, a = task_metric.plot([v[task_name] for v in val], ax=ax)\n            else:\n                raise TypeError(\n                    \"Expected argument `val` to be None or of type Dict or Sequence[Dict]. \"\n                    f\"Found type(val)= {type(val)}\"\n                )\n            fig_axs.append((f, a))\n        return fig_axs\n", "path": "src/torchmetrics/wrappers/multitask.py"}], "after_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# this is just a bypass for this module name collision with built-in one\nfrom copy import deepcopy\nfrom typing import Any, Dict, Iterable, Optional, Sequence, Tuple, Union\n\nfrom torch import Tensor, nn\n\nfrom torchmetrics.collections import MetricCollection\nfrom torchmetrics.metric import Metric\nfrom torchmetrics.utilities.imports import _MATPLOTLIB_AVAILABLE\nfrom torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE\nfrom torchmetrics.wrappers.abstract import WrapperMetric\n\nif not _MATPLOTLIB_AVAILABLE:\n    __doctest_skip__ = [\"MultitaskWrapper.plot\"]\n\n\nclass MultitaskWrapper(WrapperMetric):\n    \"\"\"Wrapper class for computing different metrics on different tasks in the context of multitask learning.\n\n    In multitask learning the different tasks requires different metrics to be evaluated. This wrapper allows\n    for easy evaluation in such cases by supporting multiple predictions and targets through a dictionary.\n    Note that only metrics where the signature of `update` follows the standard `preds, target` is supported.\n\n    Args:\n        task_metrics:\n            Dictionary associating each task to a Metric or a MetricCollection. The keys of the dictionary represent the\n            names of the tasks, and the values represent the metrics to use for each task.\n\n    Raises:\n        TypeError:\n            If argument `task_metrics` is not an dictionary\n        TypeError:\n            If not all values in the `task_metrics` dictionary is instances of `Metric` or `MetricCollection`\n\n    Example (with a single metric per class):\n         >>> import torch\n         >>> from torchmetrics.wrappers import MultitaskWrapper\n         >>> from torchmetrics.regression import MeanSquaredError\n         >>> from torchmetrics.classification import BinaryAccuracy\n         >>>\n         >>> classification_target = torch.tensor([0, 1, 0])\n         >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n         >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n         >>>\n         >>> classification_preds = torch.tensor([0, 0, 1])\n         >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n         >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n         >>>\n         >>> metrics = MultitaskWrapper({\n         ...     \"Classification\": BinaryAccuracy(),\n         ...     \"Regression\": MeanSquaredError()\n         ... })\n         >>> metrics.update(preds, targets)\n         >>> metrics.compute()\n         {'Classification': tensor(0.3333), 'Regression': tensor(0.8333)}\n\n    Example (with several metrics per task):\n         >>> import torch\n         >>> from torchmetrics import MetricCollection\n         >>> from torchmetrics.wrappers import MultitaskWrapper\n         >>> from torchmetrics.regression import MeanSquaredError, MeanAbsoluteError\n         >>> from torchmetrics.classification import BinaryAccuracy, BinaryF1Score\n         >>>\n         >>> classification_target = torch.tensor([0, 1, 0])\n         >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n         >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n         >>>\n         >>> classification_preds = torch.tensor([0, 0, 1])\n         >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n         >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n         >>>\n         >>> metrics = MultitaskWrapper({\n         ...     \"Classification\": MetricCollection(BinaryAccuracy(), BinaryF1Score()),\n         ...     \"Regression\": MetricCollection(MeanSquaredError(), MeanAbsoluteError())\n         ... })\n         >>> metrics.update(preds, targets)\n         >>> metrics.compute()\n         {'Classification': {'BinaryAccuracy': tensor(0.3333), 'BinaryF1Score': tensor(0.)},\n          'Regression': {'MeanSquaredError': tensor(0.8333), 'MeanAbsoluteError': tensor(0.6667)}}\n\n    \"\"\"\n\n    is_differentiable = False\n\n    def __init__(\n        self,\n        task_metrics: Dict[str, Union[Metric, MetricCollection]],\n    ) -> None:\n        self._check_task_metrics_type(task_metrics)\n        super().__init__()\n        self.task_metrics = nn.ModuleDict(task_metrics)\n\n    def items(self) -> Iterable[Tuple[str, nn.Module]]:\n        \"\"\"Iterate over task and task metrics.\"\"\"\n        return self.task_metrics.items()\n\n    def keys(self) -> Iterable[str]:\n        \"\"\"Iterate over task names.\"\"\"\n        return self.task_metrics.keys()\n\n    def values(self) -> Iterable[nn.Module]:\n        \"\"\"Iterate over task metrics.\"\"\"\n        return self.task_metrics.values()\n\n    @staticmethod\n    def _check_task_metrics_type(task_metrics: Dict[str, Union[Metric, MetricCollection]]) -> None:\n        if not isinstance(task_metrics, dict):\n            raise TypeError(f\"Expected argument `task_metrics` to be a dict. Found task_metrics = {task_metrics}\")\n\n        for metric in task_metrics.values():\n            if not (isinstance(metric, (Metric, MetricCollection))):\n                raise TypeError(\n                    \"Expected each task's metric to be a Metric or a MetricCollection. \"\n                    f\"Found a metric of type {type(metric)}\"\n                )\n\n    def update(self, task_preds: Dict[str, Tensor], task_targets: Dict[str, Tensor]) -> None:\n        \"\"\"Update each task's metric with its corresponding pred and target.\n\n        Args:\n            task_preds: Dictionary associating each task to a Tensor of pred.\n            task_targets: Dictionary associating each task to a Tensor of target.\n\n        \"\"\"\n        if not self.task_metrics.keys() == task_preds.keys() == task_targets.keys():\n            raise ValueError(\n                \"Expected arguments `task_preds` and `task_targets` to have the same keys as the wrapped `task_metrics`\"\n                f\". Found task_preds.keys() = {task_preds.keys()}, task_targets.keys() = {task_targets.keys()} \"\n                f\"and self.task_metrics.keys() = {self.task_metrics.keys()}\"\n            )\n\n        for task_name, metric in self.task_metrics.items():\n            pred = task_preds[task_name]\n            target = task_targets[task_name]\n            metric.update(pred, target)\n\n    def compute(self) -> Dict[str, Any]:\n        \"\"\"Compute metrics for all tasks.\"\"\"\n        return {task_name: metric.compute() for task_name, metric in self.task_metrics.items()}\n\n    def forward(self, task_preds: Dict[str, Tensor], task_targets: Dict[str, Tensor]) -> Dict[str, Any]:\n        \"\"\"Call underlying forward methods for all tasks and return the result as a dictionary.\"\"\"\n        # This method is overridden because we do not need the complex version defined in Metric, that relies on the\n        # value of full_state_update, and that also accumulates the results. Here, all computations are handled by the\n        # underlying metrics, which all have their own value of full_state_update, and which all accumulate the results\n        # by themselves.\n        return {\n            task_name: metric(task_preds[task_name], task_targets[task_name])\n            for task_name, metric in self.task_metrics.items()\n        }\n\n    def reset(self) -> None:\n        \"\"\"Reset all underlying metrics.\"\"\"\n        for metric in self.task_metrics.values():\n            metric.reset()\n        super().reset()\n\n    @staticmethod\n    def _check_arg(arg: Optional[str], name: str) -> Optional[str]:\n        if arg is None or isinstance(arg, str):\n            return arg\n        raise ValueError(f\"Expected input `{name}` to be a string, but got {type(arg)}\")\n\n    def clone(self, prefix: Optional[str] = None, postfix: Optional[str] = None) -> \"MultitaskWrapper\":\n        \"\"\"Make a copy of the metric.\n\n        Args:\n            prefix: a string to append in front of the metric keys\n            postfix: a string to append after the keys of the output dict.\n\n        \"\"\"\n        multitask_copy = deepcopy(self)\n        if prefix is not None:\n            prefix = self._check_arg(prefix, \"prefix\")\n            multitask_copy.task_metrics = nn.ModuleDict(\n                {prefix + key: value for key, value in multitask_copy.task_metrics.items()}\n            )\n        if postfix is not None:\n            postfix = self._check_arg(postfix, \"postfix\")\n            multitask_copy.task_metrics = nn.ModuleDict(\n                {key + postfix: value for key, value in multitask_copy.task_metrics.items()}\n            )\n        return multitask_copy\n\n    def plot(\n        self, val: Optional[Union[Dict, Sequence[Dict]]] = None, axes: Optional[Sequence[_AX_TYPE]] = None\n    ) -> Sequence[_PLOT_OUT_TYPE]:\n        \"\"\"Plot a single or multiple values from the metric.\n\n        All tasks' results are plotted on individual axes.\n\n        Args:\n            val: Either a single result from calling `metric.forward` or `metric.compute` or a list of these results.\n                If no value is provided, will automatically call `metric.compute` and plot that result.\n            axes: Sequence of matplotlib axis objects. If provided, will add the plots to the provided axis objects.\n                If not provided, will create them.\n\n        Returns:\n            Sequence of tuples with Figure and Axes object for each task.\n\n        .. plot::\n            :scale: 75\n\n            >>> # Example plotting a single value\n            >>> import torch\n            >>> from torchmetrics.wrappers import MultitaskWrapper\n            >>> from torchmetrics.regression import MeanSquaredError\n            >>> from torchmetrics.classification import BinaryAccuracy\n            >>>\n            >>> classification_target = torch.tensor([0, 1, 0])\n            >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n            >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n            >>>\n            >>> classification_preds = torch.tensor([0, 0, 1])\n            >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n            >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n            >>>\n            >>> metrics = MultitaskWrapper({\n            ...     \"Classification\": BinaryAccuracy(),\n            ...     \"Regression\": MeanSquaredError()\n            ... })\n            >>> metrics.update(preds, targets)\n            >>> value = metrics.compute()\n            >>> fig_, ax_ = metrics.plot(value)\n\n        .. plot::\n            :scale: 75\n\n            >>> # Example plotting multiple values\n            >>> import torch\n            >>> from torchmetrics.wrappers import MultitaskWrapper\n            >>> from torchmetrics.regression import MeanSquaredError\n            >>> from torchmetrics.classification import BinaryAccuracy\n            >>>\n            >>> classification_target = torch.tensor([0, 1, 0])\n            >>> regression_target = torch.tensor([2.5, 5.0, 4.0])\n            >>> targets = {\"Classification\": classification_target, \"Regression\": regression_target}\n            >>>\n            >>> classification_preds = torch.tensor([0, 0, 1])\n            >>> regression_preds = torch.tensor([3.0, 5.0, 2.5])\n            >>> preds = {\"Classification\": classification_preds, \"Regression\": regression_preds}\n            >>>\n            >>> metrics = MultitaskWrapper({\n            ...     \"Classification\": BinaryAccuracy(),\n            ...     \"Regression\": MeanSquaredError()\n            ... })\n            >>> values = []\n            >>> for _ in range(10):\n            ...     values.append(metrics(preds, targets))\n            >>> fig_, ax_ = metrics.plot(values)\n\n        \"\"\"\n        if axes is not None:\n            if not isinstance(axes, Sequence):\n                raise TypeError(f\"Expected argument `axes` to be a Sequence. Found type(axes) = {type(axes)}\")\n\n            if not all(isinstance(ax, _AX_TYPE) for ax in axes):\n                raise TypeError(\"Expected each ax in argument `axes` to be a matplotlib axis object\")\n\n            if len(axes) != len(self.task_metrics):\n                raise ValueError(\n                    \"Expected argument `axes` to be a Sequence of the same length as the number of tasks.\"\n                    f\"Found len(axes) = {len(axes)} and {len(self.task_metrics)} tasks\"\n                )\n\n        val = val if val is not None else self.compute()\n        fig_axs = []\n        for i, (task_name, task_metric) in enumerate(self.task_metrics.items()):\n            ax = axes[i] if axes is not None else None\n            if isinstance(val, Dict):\n                f, a = task_metric.plot(val[task_name], ax=ax)\n            elif isinstance(val, Sequence):\n                f, a = task_metric.plot([v[task_name] for v in val], ax=ax)\n            else:\n                raise TypeError(\n                    \"Expected argument `val` to be None or of type Dict or Sequence[Dict]. \"\n                    f\"Found type(val)= {type(val)}\"\n                )\n            fig_axs.append((f, a))\n        return fig_axs\n", "path": "src/torchmetrics/wrappers/multitask.py"}]}