Prefix postfix args in clone for multitask (#2330)

* feature
* tests
* docstring

Author: SkafteNicki

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 @@ def reset(self) -> None:
             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]:

tests/unittests/wrappers/test_multitask.py

@@ -207,3 +207,18 @@ def test_nested_multitask_wrapper():
     multitask_results = multitask_metrics.compute()
 
     assert _dict_results_same_as_individual_results(classification_results, regression_results, multitask_results)
+
+
+def test_clone_with_prefix_and_postfix():
+    """Check that the clone method works with prefix and postfix arguments."""
+    multitask_metrics = MultitaskWrapper({"Classification": BinaryAccuracy(), "Regression": MeanSquaredError()})
+    cloned_metrics_with_prefix = multitask_metrics.clone(prefix="prefix_")
+    cloned_metrics_with_postfix = multitask_metrics.clone(postfix="_postfix")
+
+    # Check if the cloned metrics have the expected keys
+    assert set(cloned_metrics_with_prefix.task_metrics.keys()) == {"prefix_Classification", "prefix_Regression"}
+    assert set(cloned_metrics_with_postfix.task_metrics.keys()) == {"Classification_postfix", "Regression_postfix"}
+
+    # Check if the cloned metrics have the expected values
+    assert isinstance(cloned_metrics_with_prefix.task_metrics["prefix_Classification"], BinaryAccuracy)
+    assert isinstance(cloned_metrics_with_prefix.task_metrics["prefix_Regression"], MeanSquaredError)