We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Half precision doesn't work with top_k > 1
## 🐛 Bug

When using predictions at half-precision, MulticlassAccuracy throws an error when using top_k > 1. Please see the code below:

### To Reproduce

Steps to reproduce the behavior...

```py
import torch
from torchmetrics.classification import MulticlassAccuracy

preds = torch.tensor([[1, 0, 0]]).half()
targets = torch.tensor([1])

metric = MulticlassAccuracy(num_classes=3, top_k=2)
metric(preds, targets)
metric.compute()
```


### Environment

- TorchMetrics version: 1.0.1
- Python & PyTorch Version (e.g., 1.0): 3.10 and 2.0.1



--- 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/utilities/data.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 import sys
15 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
16 
17 import torch
18 from lightning_utilities import apply_to_collection
19 from torch import Tensor
20 
21 from torchmetrics.utilities.exceptions import TorchMetricsUserWarning
22 from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12, _XLA_AVAILABLE
23 from torchmetrics.utilities.prints import rank_zero_warn
24 
25 METRIC_EPS = 1e-6
26 
27 
28 def dim_zero_cat(x: Union[Tensor, List[Tensor]]) -> Tensor:
29     """Concatenation along the zero dimension."""
30     if isinstance(x, torch.Tensor):
31         return x
32     x = [y.unsqueeze(0) if y.numel() == 1 and y.ndim == 0 else y for y in x]
33     if not x:  # empty list
34         raise ValueError("No samples to concatenate")
35     return torch.cat(x, dim=0)
36 
37 
38 def dim_zero_sum(x: Tensor) -> Tensor:
39     """Summation along the zero dimension."""
40     return torch.sum(x, dim=0)
41 
42 
43 def dim_zero_mean(x: Tensor) -> Tensor:
44     """Average along the zero dimension."""
45     return torch.mean(x, dim=0)
46 
47 
48 def dim_zero_max(x: Tensor) -> Tensor:
49     """Max along the zero dimension."""
50     return torch.max(x, dim=0).values
51 
52 
53 def dim_zero_min(x: Tensor) -> Tensor:
54     """Min along the zero dimension."""
55     return torch.min(x, dim=0).values
56 
57 
58 def _flatten(x: Sequence) -> list:
59     """Flatten list of list into single list."""
60     return [item for sublist in x for item in sublist]
61 
62 
63 def _flatten_dict(x: Dict) -> Tuple[Dict, bool]:
64     """Flatten dict of dicts into single dict and checking for duplicates in keys along the way."""
65     new_dict = {}
66     duplicates = False
67     for key, value in x.items():
68         if isinstance(value, dict):
69             for k, v in value.items():
70                 if k in new_dict:
71                     duplicates = True
72                 new_dict[k] = v
73         else:
74             if key in new_dict:
75                 duplicates = True
76             new_dict[key] = value
77     return new_dict, duplicates
78 
79 
80 def to_onehot(
81     label_tensor: Tensor,
82     num_classes: Optional[int] = None,
83 ) -> Tensor:
84     """Convert  a dense label tensor to one-hot format.
85 
86     Args:
87         label_tensor: dense label tensor, with shape [N, d1, d2, ...]
88         num_classes: number of classes C
89 
90     Returns:
91         A sparse label tensor with shape [N, C, d1, d2, ...]
92 
93     Example:
94         >>> x = torch.tensor([1, 2, 3])
95         >>> to_onehot(x)
96         tensor([[0, 1, 0, 0],
97                 [0, 0, 1, 0],
98                 [0, 0, 0, 1]])
99 
100     """
101     if num_classes is None:
102         num_classes = int(label_tensor.max().detach().item() + 1)
103 
104     tensor_onehot = torch.zeros(
105         label_tensor.shape[0],
106         num_classes,
107         *label_tensor.shape[1:],
108         dtype=label_tensor.dtype,
109         device=label_tensor.device,
110     )
111     index = label_tensor.long().unsqueeze(1).expand_as(tensor_onehot)
112     return tensor_onehot.scatter_(1, index, 1.0)
113 
114 
115 def select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:
116     """Convert a probability tensor to binary by selecting top-k the highest entries.
117 
118     Args:
119         prob_tensor: dense tensor of shape ``[..., C, ...]``, where ``C`` is in the
120             position defined by the ``dim`` argument
121         topk: number of the highest entries to turn into 1s
122         dim: dimension on which to compare entries
123 
124     Returns:
125         A binary tensor of the same shape as the input tensor of type ``torch.int32``
126 
127     Example:
128         >>> x = torch.tensor([[1.1, 2.0, 3.0], [2.0, 1.0, 0.5]])
129         >>> select_topk(x, topk=2)
130         tensor([[0, 1, 1],
131                 [1, 1, 0]], dtype=torch.int32)
132 
133     """
134     zeros = torch.zeros_like(prob_tensor)
135     if topk == 1:  # argmax has better performance than topk
136         topk_tensor = zeros.scatter(dim, prob_tensor.argmax(dim=dim, keepdim=True), 1.0)
137     else:
138         topk_tensor = zeros.scatter(dim, prob_tensor.topk(k=topk, dim=dim).indices, 1.0)
139     return topk_tensor.int()
140 
141 
142 def to_categorical(x: Tensor, argmax_dim: int = 1) -> Tensor:
143     """Convert  a tensor of probabilities to a dense label tensor.
144 
145     Args:
146         x: probabilities to get the categorical label [N, d1, d2, ...]
147         argmax_dim: dimension to apply
148 
149     Return:
150         A tensor with categorical labels [N, d2, ...]
151 
152     Example:
153         >>> x = torch.tensor([[0.2, 0.5], [0.9, 0.1]])
154         >>> to_categorical(x)
155         tensor([1, 0])
156 
157     """
158     return torch.argmax(x, dim=argmax_dim)
159 
160 
161 def _squeeze_scalar_element_tensor(x: Tensor) -> Tensor:
162     return x.squeeze() if x.numel() == 1 else x
163 
164 
165 def _squeeze_if_scalar(data: Any) -> Any:
166     return apply_to_collection(data, Tensor, _squeeze_scalar_element_tensor)
167 
168 
169 def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:
170     """Implement custom bincount.
171 
172     PyTorch currently does not support ``torch.bincount`` when running in deterministic mode on GPU or when running
173     MPS devices or when running on XLA device. This implementation therefore falls back to using a combination of
174     `torch.arange` and `torch.eq` in these scenarios. A small performance hit can expected and higher memory consumption
175     as `[batch_size, mincount]` tensor needs to be initialized compared to native ``torch.bincount``.
176 
177     Args:
178         x: tensor to count
179         minlength: minimum length to count
180 
181     Returns:
182         Number of occurrences for each unique element in x
183 
184     Example:
185         >>> x = torch.tensor([0,0,0,1,1,2,2,2,2])
186         >>> _bincount(x, minlength=3)
187         tensor([3, 2, 4])
188 
189     """
190     if minlength is None:
191         minlength = len(torch.unique(x))
192 
193     if torch.are_deterministic_algorithms_enabled() or _XLA_AVAILABLE or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:
194         mesh = torch.arange(minlength, device=x.device).repeat(len(x), 1)
195         return torch.eq(x.reshape(-1, 1), mesh).sum(dim=0)
196 
197     return torch.bincount(x, minlength=minlength)
198 
199 
200 def _cumsum(x: Tensor, dim: Optional[int] = 0, dtype: Optional[torch.dtype] = None) -> Tensor:
201     if torch.are_deterministic_algorithms_enabled() and x.is_cuda and x.is_floating_point() and sys.platform != "win32":
202         rank_zero_warn(
203             "You are trying to use a metric in deterministic mode on GPU that uses `torch.cumsum`, which is currently "
204             "not supported. The tensor will be copied to the CPU memory to compute it and then copied back to GPU. "
205             "Expect some slowdowns.",
206             TorchMetricsUserWarning,
207         )
208         return x.cpu().cumsum(dim=dim, dtype=dtype).cuda()
209     return torch.cumsum(x, dim=dim, dtype=dtype)
210 
211 
212 def _flexible_bincount(x: Tensor) -> Tensor:
213     """Similar to `_bincount`, but works also with tensor that do not contain continuous values.
214 
215     Args:
216         x: tensor to count
217 
218     Returns:
219         Number of occurrences for each unique element in x
220 
221     """
222     # make sure elements in x start from 0
223     x = x - x.min()
224     unique_x = torch.unique(x)
225 
226     output = _bincount(x, minlength=torch.max(unique_x) + 1)  # type: ignore[arg-type]
227     # remove zeros from output tensor
228     return output[unique_x]
229 
230 
231 def allclose(tensor1: Tensor, tensor2: Tensor) -> bool:
232     """Wrap torch.allclose to be robust towards dtype difference."""
233     if tensor1.dtype != tensor2.dtype:
234         tensor2 = tensor2.to(dtype=tensor1.dtype)
235     return torch.allclose(tensor1, tensor2)
236
```
--- 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/utilities/data.py b/src/torchmetrics/utilities/data.py
--- a/src/torchmetrics/utilities/data.py
+++ b/src/torchmetrics/utilities/data.py
@@ -112,6 +112,14 @@
     return tensor_onehot.scatter_(1, index, 1.0)
 
 
+def _top_k_with_half_precision_support(x: Tensor, k: int = 1, dim: int = 1) -> Tensor:
+    """torch.top_k does not support half precision on CPU."""
+    if x.dtype == torch.half and not x.is_cuda:
+        idx = torch.argsort(x, dim=dim, descending=True)
+        return idx.narrow(dim, 0, k)
+    return x.topk(k=k, dim=dim).indices
+
+
 def select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:
     """Convert a probability tensor to binary by selecting top-k the highest entries.
 
@@ -135,7 +143,7 @@
     if topk == 1:  # argmax has better performance than topk
         topk_tensor = zeros.scatter(dim, prob_tensor.argmax(dim=dim, keepdim=True), 1.0)
     else:
-        topk_tensor = zeros.scatter(dim, prob_tensor.topk(k=topk, dim=dim).indices, 1.0)
+        topk_tensor = zeros.scatter(dim, _top_k_with_half_precision_support(prob_tensor, k=topk, dim=dim), 1.0)
     return topk_tensor.int()

{"golden_diff": "diff --git a/src/torchmetrics/utilities/data.py b/src/torchmetrics/utilities/data.py\n--- a/src/torchmetrics/utilities/data.py\n+++ b/src/torchmetrics/utilities/data.py\n@@ -112,6 +112,14 @@\n     return tensor_onehot.scatter_(1, index, 1.0)\n \n \n+def _top_k_with_half_precision_support(x: Tensor, k: int = 1, dim: int = 1) -> Tensor:\n+    \"\"\"torch.top_k does not support half precision on CPU.\"\"\"\n+    if x.dtype == torch.half and not x.is_cuda:\n+        idx = torch.argsort(x, dim=dim, descending=True)\n+        return idx.narrow(dim, 0, k)\n+    return x.topk(k=k, dim=dim).indices\n+\n+\n def select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:\n     \"\"\"Convert a probability tensor to binary by selecting top-k the highest entries.\n \n@@ -135,7 +143,7 @@\n     if topk == 1:  # argmax has better performance than topk\n         topk_tensor = zeros.scatter(dim, prob_tensor.argmax(dim=dim, keepdim=True), 1.0)\n     else:\n-        topk_tensor = zeros.scatter(dim, prob_tensor.topk(k=topk, dim=dim).indices, 1.0)\n+        topk_tensor = zeros.scatter(dim, _top_k_with_half_precision_support(prob_tensor, k=topk, dim=dim), 1.0)\n     return topk_tensor.int()\n", "issue": "Half precision doesn't work with top_k > 1\n## \ud83d\udc1b Bug\r\n\r\nWhen using predictions at half-precision, MulticlassAccuracy throws an error when using top_k > 1. Please see the code below:\r\n\r\n### To Reproduce\r\n\r\nSteps to reproduce the behavior...\r\n\r\n```py\r\nimport torch\r\nfrom torchmetrics.classification import MulticlassAccuracy\r\n\r\npreds = torch.tensor([[1, 0, 0]]).half()\r\ntargets = torch.tensor([1])\r\n\r\nmetric = MulticlassAccuracy(num_classes=3, top_k=2)\r\nmetric(preds, targets)\r\nmetric.compute()\r\n```\r\n\r\n\r\n### Environment\r\n\r\n- TorchMetrics version: 1.0.1\r\n- Python & PyTorch Version (e.g., 1.0): 3.10 and 2.0.1\r\n\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.\nimport sys\nfrom typing import Any, Dict, List, Optional, Sequence, Tuple, Union\n\nimport torch\nfrom lightning_utilities import apply_to_collection\nfrom torch import Tensor\n\nfrom torchmetrics.utilities.exceptions import TorchMetricsUserWarning\nfrom torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12, _XLA_AVAILABLE\nfrom torchmetrics.utilities.prints import rank_zero_warn\n\nMETRIC_EPS = 1e-6\n\n\ndef dim_zero_cat(x: Union[Tensor, List[Tensor]]) -> Tensor:\n    \"\"\"Concatenation along the zero dimension.\"\"\"\n    if isinstance(x, torch.Tensor):\n        return x\n    x = [y.unsqueeze(0) if y.numel() == 1 and y.ndim == 0 else y for y in x]\n    if not x:  # empty list\n        raise ValueError(\"No samples to concatenate\")\n    return torch.cat(x, dim=0)\n\n\ndef dim_zero_sum(x: Tensor) -> Tensor:\n    \"\"\"Summation along the zero dimension.\"\"\"\n    return torch.sum(x, dim=0)\n\n\ndef dim_zero_mean(x: Tensor) -> Tensor:\n    \"\"\"Average along the zero dimension.\"\"\"\n    return torch.mean(x, dim=0)\n\n\ndef dim_zero_max(x: Tensor) -> Tensor:\n    \"\"\"Max along the zero dimension.\"\"\"\n    return torch.max(x, dim=0).values\n\n\ndef dim_zero_min(x: Tensor) -> Tensor:\n    \"\"\"Min along the zero dimension.\"\"\"\n    return torch.min(x, dim=0).values\n\n\ndef _flatten(x: Sequence) -> list:\n    \"\"\"Flatten list of list into single list.\"\"\"\n    return [item for sublist in x for item in sublist]\n\n\ndef _flatten_dict(x: Dict) -> Tuple[Dict, bool]:\n    \"\"\"Flatten dict of dicts into single dict and checking for duplicates in keys along the way.\"\"\"\n    new_dict = {}\n    duplicates = False\n    for key, value in x.items():\n        if isinstance(value, dict):\n            for k, v in value.items():\n                if k in new_dict:\n                    duplicates = True\n                new_dict[k] = v\n        else:\n            if key in new_dict:\n                duplicates = True\n            new_dict[key] = value\n    return new_dict, duplicates\n\n\ndef to_onehot(\n    label_tensor: Tensor,\n    num_classes: Optional[int] = None,\n) -> Tensor:\n    \"\"\"Convert  a dense label tensor to one-hot format.\n\n    Args:\n        label_tensor: dense label tensor, with shape [N, d1, d2, ...]\n        num_classes: number of classes C\n\n    Returns:\n        A sparse label tensor with shape [N, C, d1, d2, ...]\n\n    Example:\n        >>> x = torch.tensor([1, 2, 3])\n        >>> to_onehot(x)\n        tensor([[0, 1, 0, 0],\n                [0, 0, 1, 0],\n                [0, 0, 0, 1]])\n\n    \"\"\"\n    if num_classes is None:\n        num_classes = int(label_tensor.max().detach().item() + 1)\n\n    tensor_onehot = torch.zeros(\n        label_tensor.shape[0],\n        num_classes,\n        *label_tensor.shape[1:],\n        dtype=label_tensor.dtype,\n        device=label_tensor.device,\n    )\n    index = label_tensor.long().unsqueeze(1).expand_as(tensor_onehot)\n    return tensor_onehot.scatter_(1, index, 1.0)\n\n\ndef select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:\n    \"\"\"Convert a probability tensor to binary by selecting top-k the highest entries.\n\n    Args:\n        prob_tensor: dense tensor of shape ``[..., C, ...]``, where ``C`` is in the\n            position defined by the ``dim`` argument\n        topk: number of the highest entries to turn into 1s\n        dim: dimension on which to compare entries\n\n    Returns:\n        A binary tensor of the same shape as the input tensor of type ``torch.int32``\n\n    Example:\n        >>> x = torch.tensor([[1.1, 2.0, 3.0], [2.0, 1.0, 0.5]])\n        >>> select_topk(x, topk=2)\n        tensor([[0, 1, 1],\n                [1, 1, 0]], dtype=torch.int32)\n\n    \"\"\"\n    zeros = torch.zeros_like(prob_tensor)\n    if topk == 1:  # argmax has better performance than topk\n        topk_tensor = zeros.scatter(dim, prob_tensor.argmax(dim=dim, keepdim=True), 1.0)\n    else:\n        topk_tensor = zeros.scatter(dim, prob_tensor.topk(k=topk, dim=dim).indices, 1.0)\n    return topk_tensor.int()\n\n\ndef to_categorical(x: Tensor, argmax_dim: int = 1) -> Tensor:\n    \"\"\"Convert  a tensor of probabilities to a dense label tensor.\n\n    Args:\n        x: probabilities to get the categorical label [N, d1, d2, ...]\n        argmax_dim: dimension to apply\n\n    Return:\n        A tensor with categorical labels [N, d2, ...]\n\n    Example:\n        >>> x = torch.tensor([[0.2, 0.5], [0.9, 0.1]])\n        >>> to_categorical(x)\n        tensor([1, 0])\n\n    \"\"\"\n    return torch.argmax(x, dim=argmax_dim)\n\n\ndef _squeeze_scalar_element_tensor(x: Tensor) -> Tensor:\n    return x.squeeze() if x.numel() == 1 else x\n\n\ndef _squeeze_if_scalar(data: Any) -> Any:\n    return apply_to_collection(data, Tensor, _squeeze_scalar_element_tensor)\n\n\ndef _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:\n    \"\"\"Implement custom bincount.\n\n    PyTorch currently does not support ``torch.bincount`` when running in deterministic mode on GPU or when running\n    MPS devices or when running on XLA device. This implementation therefore falls back to using a combination of\n    `torch.arange` and `torch.eq` in these scenarios. A small performance hit can expected and higher memory consumption\n    as `[batch_size, mincount]` tensor needs to be initialized compared to native ``torch.bincount``.\n\n    Args:\n        x: tensor to count\n        minlength: minimum length to count\n\n    Returns:\n        Number of occurrences for each unique element in x\n\n    Example:\n        >>> x = torch.tensor([0,0,0,1,1,2,2,2,2])\n        >>> _bincount(x, minlength=3)\n        tensor([3, 2, 4])\n\n    \"\"\"\n    if minlength is None:\n        minlength = len(torch.unique(x))\n\n    if torch.are_deterministic_algorithms_enabled() or _XLA_AVAILABLE or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:\n        mesh = torch.arange(minlength, device=x.device).repeat(len(x), 1)\n        return torch.eq(x.reshape(-1, 1), mesh).sum(dim=0)\n\n    return torch.bincount(x, minlength=minlength)\n\n\ndef _cumsum(x: Tensor, dim: Optional[int] = 0, dtype: Optional[torch.dtype] = None) -> Tensor:\n    if torch.are_deterministic_algorithms_enabled() and x.is_cuda and x.is_floating_point() and sys.platform != \"win32\":\n        rank_zero_warn(\n            \"You are trying to use a metric in deterministic mode on GPU that uses `torch.cumsum`, which is currently \"\n            \"not supported. The tensor will be copied to the CPU memory to compute it and then copied back to GPU. \"\n            \"Expect some slowdowns.\",\n            TorchMetricsUserWarning,\n        )\n        return x.cpu().cumsum(dim=dim, dtype=dtype).cuda()\n    return torch.cumsum(x, dim=dim, dtype=dtype)\n\n\ndef _flexible_bincount(x: Tensor) -> Tensor:\n    \"\"\"Similar to `_bincount`, but works also with tensor that do not contain continuous values.\n\n    Args:\n        x: tensor to count\n\n    Returns:\n        Number of occurrences for each unique element in x\n\n    \"\"\"\n    # make sure elements in x start from 0\n    x = x - x.min()\n    unique_x = torch.unique(x)\n\n    output = _bincount(x, minlength=torch.max(unique_x) + 1)  # type: ignore[arg-type]\n    # remove zeros from output tensor\n    return output[unique_x]\n\n\ndef allclose(tensor1: Tensor, tensor2: Tensor) -> bool:\n    \"\"\"Wrap torch.allclose to be robust towards dtype difference.\"\"\"\n    if tensor1.dtype != tensor2.dtype:\n        tensor2 = tensor2.to(dtype=tensor1.dtype)\n    return torch.allclose(tensor1, tensor2)\n", "path": "src/torchmetrics/utilities/data.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.\nimport sys\nfrom typing import Any, Dict, List, Optional, Sequence, Tuple, Union\n\nimport torch\nfrom lightning_utilities import apply_to_collection\nfrom torch import Tensor\n\nfrom torchmetrics.utilities.exceptions import TorchMetricsUserWarning\nfrom torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12, _XLA_AVAILABLE\nfrom torchmetrics.utilities.prints import rank_zero_warn\n\nMETRIC_EPS = 1e-6\n\n\ndef dim_zero_cat(x: Union[Tensor, List[Tensor]]) -> Tensor:\n    \"\"\"Concatenation along the zero dimension.\"\"\"\n    if isinstance(x, torch.Tensor):\n        return x\n    x = [y.unsqueeze(0) if y.numel() == 1 and y.ndim == 0 else y for y in x]\n    if not x:  # empty list\n        raise ValueError(\"No samples to concatenate\")\n    return torch.cat(x, dim=0)\n\n\ndef dim_zero_sum(x: Tensor) -> Tensor:\n    \"\"\"Summation along the zero dimension.\"\"\"\n    return torch.sum(x, dim=0)\n\n\ndef dim_zero_mean(x: Tensor) -> Tensor:\n    \"\"\"Average along the zero dimension.\"\"\"\n    return torch.mean(x, dim=0)\n\n\ndef dim_zero_max(x: Tensor) -> Tensor:\n    \"\"\"Max along the zero dimension.\"\"\"\n    return torch.max(x, dim=0).values\n\n\ndef dim_zero_min(x: Tensor) -> Tensor:\n    \"\"\"Min along the zero dimension.\"\"\"\n    return torch.min(x, dim=0).values\n\n\ndef _flatten(x: Sequence) -> list:\n    \"\"\"Flatten list of list into single list.\"\"\"\n    return [item for sublist in x for item in sublist]\n\n\ndef _flatten_dict(x: Dict) -> Tuple[Dict, bool]:\n    \"\"\"Flatten dict of dicts into single dict and checking for duplicates in keys along the way.\"\"\"\n    new_dict = {}\n    duplicates = False\n    for key, value in x.items():\n        if isinstance(value, dict):\n            for k, v in value.items():\n                if k in new_dict:\n                    duplicates = True\n                new_dict[k] = v\n        else:\n            if key in new_dict:\n                duplicates = True\n            new_dict[key] = value\n    return new_dict, duplicates\n\n\ndef to_onehot(\n    label_tensor: Tensor,\n    num_classes: Optional[int] = None,\n) -> Tensor:\n    \"\"\"Convert  a dense label tensor to one-hot format.\n\n    Args:\n        label_tensor: dense label tensor, with shape [N, d1, d2, ...]\n        num_classes: number of classes C\n\n    Returns:\n        A sparse label tensor with shape [N, C, d1, d2, ...]\n\n    Example:\n        >>> x = torch.tensor([1, 2, 3])\n        >>> to_onehot(x)\n        tensor([[0, 1, 0, 0],\n                [0, 0, 1, 0],\n                [0, 0, 0, 1]])\n\n    \"\"\"\n    if num_classes is None:\n        num_classes = int(label_tensor.max().detach().item() + 1)\n\n    tensor_onehot = torch.zeros(\n        label_tensor.shape[0],\n        num_classes,\n        *label_tensor.shape[1:],\n        dtype=label_tensor.dtype,\n        device=label_tensor.device,\n    )\n    index = label_tensor.long().unsqueeze(1).expand_as(tensor_onehot)\n    return tensor_onehot.scatter_(1, index, 1.0)\n\n\ndef _top_k_with_half_precision_support(x: Tensor, k: int = 1, dim: int = 1) -> Tensor:\n    \"\"\"torch.top_k does not support half precision on CPU.\"\"\"\n    if x.dtype == torch.half and not x.is_cuda:\n        idx = torch.argsort(x, dim=dim, descending=True)\n        return idx.narrow(dim, 0, k)\n    return x.topk(k=k, dim=dim).indices\n\n\ndef select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:\n    \"\"\"Convert a probability tensor to binary by selecting top-k the highest entries.\n\n    Args:\n        prob_tensor: dense tensor of shape ``[..., C, ...]``, where ``C`` is in the\n            position defined by the ``dim`` argument\n        topk: number of the highest entries to turn into 1s\n        dim: dimension on which to compare entries\n\n    Returns:\n        A binary tensor of the same shape as the input tensor of type ``torch.int32``\n\n    Example:\n        >>> x = torch.tensor([[1.1, 2.0, 3.0], [2.0, 1.0, 0.5]])\n        >>> select_topk(x, topk=2)\n        tensor([[0, 1, 1],\n                [1, 1, 0]], dtype=torch.int32)\n\n    \"\"\"\n    zeros = torch.zeros_like(prob_tensor)\n    if topk == 1:  # argmax has better performance than topk\n        topk_tensor = zeros.scatter(dim, prob_tensor.argmax(dim=dim, keepdim=True), 1.0)\n    else:\n        topk_tensor = zeros.scatter(dim, _top_k_with_half_precision_support(prob_tensor, k=topk, dim=dim), 1.0)\n    return topk_tensor.int()\n\n\ndef to_categorical(x: Tensor, argmax_dim: int = 1) -> Tensor:\n    \"\"\"Convert  a tensor of probabilities to a dense label tensor.\n\n    Args:\n        x: probabilities to get the categorical label [N, d1, d2, ...]\n        argmax_dim: dimension to apply\n\n    Return:\n        A tensor with categorical labels [N, d2, ...]\n\n    Example:\n        >>> x = torch.tensor([[0.2, 0.5], [0.9, 0.1]])\n        >>> to_categorical(x)\n        tensor([1, 0])\n\n    \"\"\"\n    return torch.argmax(x, dim=argmax_dim)\n\n\ndef _squeeze_scalar_element_tensor(x: Tensor) -> Tensor:\n    return x.squeeze() if x.numel() == 1 else x\n\n\ndef _squeeze_if_scalar(data: Any) -> Any:\n    return apply_to_collection(data, Tensor, _squeeze_scalar_element_tensor)\n\n\ndef _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:\n    \"\"\"Implement custom bincount.\n\n    PyTorch currently does not support ``torch.bincount`` when running in deterministic mode on GPU or when running\n    MPS devices or when running on XLA device. This implementation therefore falls back to using a combination of\n    `torch.arange` and `torch.eq` in these scenarios. A small performance hit can expected and higher memory consumption\n    as `[batch_size, mincount]` tensor needs to be initialized compared to native ``torch.bincount``.\n\n    Args:\n        x: tensor to count\n        minlength: minimum length to count\n\n    Returns:\n        Number of occurrences for each unique element in x\n\n    Example:\n        >>> x = torch.tensor([0,0,0,1,1,2,2,2,2])\n        >>> _bincount(x, minlength=3)\n        tensor([3, 2, 4])\n\n    \"\"\"\n    if minlength is None:\n        minlength = len(torch.unique(x))\n\n    if torch.are_deterministic_algorithms_enabled() or _XLA_AVAILABLE or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:\n        mesh = torch.arange(minlength, device=x.device).repeat(len(x), 1)\n        return torch.eq(x.reshape(-1, 1), mesh).sum(dim=0)\n\n    return torch.bincount(x, minlength=minlength)\n\n\ndef _cumsum(x: Tensor, dim: Optional[int] = 0, dtype: Optional[torch.dtype] = None) -> Tensor:\n    if torch.are_deterministic_algorithms_enabled() and x.is_cuda and x.is_floating_point() and sys.platform != \"win32\":\n        rank_zero_warn(\n            \"You are trying to use a metric in deterministic mode on GPU that uses `torch.cumsum`, which is currently \"\n            \"not supported. The tensor will be copied to the CPU memory to compute it and then copied back to GPU. \"\n            \"Expect some slowdowns.\",\n            TorchMetricsUserWarning,\n        )\n        return x.cpu().cumsum(dim=dim, dtype=dtype).cuda()\n    return torch.cumsum(x, dim=dim, dtype=dtype)\n\n\ndef _flexible_bincount(x: Tensor) -> Tensor:\n    \"\"\"Similar to `_bincount`, but works also with tensor that do not contain continuous values.\n\n    Args:\n        x: tensor to count\n\n    Returns:\n        Number of occurrences for each unique element in x\n\n    \"\"\"\n    # make sure elements in x start from 0\n    x = x - x.min()\n    unique_x = torch.unique(x)\n\n    output = _bincount(x, minlength=torch.max(unique_x) + 1)  # type: ignore[arg-type]\n    # remove zeros from output tensor\n    return output[unique_x]\n\n\ndef allclose(tensor1: Tensor, tensor2: Tensor) -> bool:\n    \"\"\"Wrap torch.allclose to be robust towards dtype difference.\"\"\"\n    if tensor1.dtype != tensor2.dtype:\n        tensor2 = tensor2.to(dtype=tensor1.dtype)\n    return torch.allclose(tensor1, tensor2)\n", "path": "src/torchmetrics/utilities/data.py"}]}