We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
MatthewsCorrCoef has massive slow down on both CPU and especially GPU
## 🐛 Bug

It's extremely slow to compute the Matthew Correlation Coefficient since torchmetrics == 0.10.0.

### To Reproduce

During my testing, I am seeing massive slow down when computing the Matthew Correlation Coefficient, especially on the GPU (using pytorch-lightning to build, train and test a deep learning model's performance on this metric).

I have compiled a code sample (next section) and below shows the result over different versions of torchmetrics. 

ON CPU:
torchmetrics 0.7.3: 0.93961 s
torchmetrics 0.8.0: 0.93549 s
torchmetrics 0.8.2: 0.94494 s
torchmetrics 0.9.0: 0.92856 s
torchmetrics 0.9.2: 0.93682 s
torchmetrics 0.10.0: 1.10903 s

ON GPU:
torchmetrics 0.7.3: 0.11444 s
torchmetrics 0.8.0: 0.11682 s
torchmetrics 0.8.2: 0.11425 s
torchmetrics 0.9.0: 0.11433 s
torchmetrics 0.9.2: 0.11410 s
torchmetrics 0.10.0: 359.30208 s

So yeah testing over thousands of batches now takes almost a week for me to complete :rofl: Please take a look at this soon.

#### Code sample

```
from tqdm.auto import tqdm
import time
import torch
from torchmetrics import MatthewsCorrCoef

torch.manual_seed(1)

b, h, w = 10, 1080, 1920
device = "cpu"

def generate(b, h, w):
    prob = torch.rand(b, h, w).to(device)
    truth = torch.randint(0, 2, (b, h, w)).to(device)
    return prob, truth

batches = []
for _ in range(10):
    batches.append(generate(b, h, w))

mcc = MatthewsCorrCoef(num_classes=2).to(device)

t1 = time.time()
for detections, targets in tqdm(batches):
    mcc.update(detections, targets)
print(f"{time.time() - t1:.5f}")
```

### Expected behavior

It's supposed to be faster in the range of ~0.1 second for GPU and ~0.9 second for CPU in this naive benchmark.

### Environment

- GPU: NVIDIA RTX 3090
- TorchMetrics version (`pip`): 0.10.0
- Python & PyTorch Version (e.g., 1.0): python=3.9, pytorch=1.12.1
- Any other relevant information such as OS (e.g., Linux): ubuntu18.04, (same behavior with 20.04)

--- 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 PyTorch 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 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union
15 
16 import torch
17 from torch import Tensor
18 
19 from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12
20 
21 METRIC_EPS = 1e-6
22 
23 
24 def dim_zero_cat(x: Union[Tensor, List[Tensor]]) -> Tensor:
25     """Concatenation along the zero dimension."""
26     x = x if isinstance(x, (list, tuple)) else [x]
27     x = [y.unsqueeze(0) if y.numel() == 1 and y.ndim == 0 else y for y in x]
28     if not x:  # empty list
29         raise ValueError("No samples to concatenate")
30     return torch.cat(x, dim=0)
31 
32 
33 def dim_zero_sum(x: Tensor) -> Tensor:
34     """Summation along the zero dimension."""
35     return torch.sum(x, dim=0)
36 
37 
38 def dim_zero_mean(x: Tensor) -> Tensor:
39     """Average along the zero dimension."""
40     return torch.mean(x, dim=0)
41 
42 
43 def dim_zero_max(x: Tensor) -> Tensor:
44     """Max along the zero dimension."""
45     return torch.max(x, dim=0).values
46 
47 
48 def dim_zero_min(x: Tensor) -> Tensor:
49     """Min along the zero dimension."""
50     return torch.min(x, dim=0).values
51 
52 
53 def _flatten(x: Sequence) -> list:
54     """Flatten list of list into single list."""
55     return [item for sublist in x for item in sublist]
56 
57 
58 def _flatten_dict(x: Dict) -> Dict:
59     """Flatten dict of dicts into single dict."""
60     new_dict = {}
61     for key, value in x.items():
62         if isinstance(value, dict):
63             for k, v in value.items():
64                 new_dict[k] = v
65         else:
66             new_dict[key] = value
67     return new_dict
68 
69 
70 def to_onehot(
71     label_tensor: Tensor,
72     num_classes: Optional[int] = None,
73 ) -> Tensor:
74     """Converts a dense label tensor to one-hot format.
75 
76     Args:
77         label_tensor: dense label tensor, with shape [N, d1, d2, ...]
78         num_classes: number of classes C
79 
80     Returns:
81         A sparse label tensor with shape [N, C, d1, d2, ...]
82 
83     Example:
84         >>> x = torch.tensor([1, 2, 3])
85         >>> to_onehot(x)
86         tensor([[0, 1, 0, 0],
87                 [0, 0, 1, 0],
88                 [0, 0, 0, 1]])
89     """
90     if num_classes is None:
91         num_classes = int(label_tensor.max().detach().item() + 1)
92 
93     tensor_onehot = torch.zeros(
94         label_tensor.shape[0],
95         num_classes,
96         *label_tensor.shape[1:],
97         dtype=label_tensor.dtype,
98         device=label_tensor.device,
99     )
100     index = label_tensor.long().unsqueeze(1).expand_as(tensor_onehot)
101     return tensor_onehot.scatter_(1, index, 1.0)
102 
103 
104 def select_topk(prob_tensor: Tensor, topk: int = 1, dim: int = 1) -> Tensor:
105     """Convert a probability tensor to binary by selecting top-k the highest entries.
106 
107     Args:
108         prob_tensor: dense tensor of shape ``[..., C, ...]``, where ``C`` is in the
109             position defined by the ``dim`` argument
110         topk: number of the highest entries to turn into 1s
111         dim: dimension on which to compare entries
112 
113     Returns:
114         A binary tensor of the same shape as the input tensor of type ``torch.int32``
115 
116     Example:
117         >>> x = torch.tensor([[1.1, 2.0, 3.0], [2.0, 1.0, 0.5]])
118         >>> select_topk(x, topk=2)
119         tensor([[0, 1, 1],
120                 [1, 1, 0]], dtype=torch.int32)
121     """
122     zeros = torch.zeros_like(prob_tensor)
123     if topk == 1:  # argmax has better performance than topk
124         topk_tensor = zeros.scatter(dim, prob_tensor.argmax(dim=dim, keepdim=True), 1.0)
125     else:
126         topk_tensor = zeros.scatter(dim, prob_tensor.topk(k=topk, dim=dim).indices, 1.0)
127     return topk_tensor.int()
128 
129 
130 def to_categorical(x: Tensor, argmax_dim: int = 1) -> Tensor:
131     """Converts a tensor of probabilities to a dense label tensor.
132 
133     Args:
134         x: probabilities to get the categorical label [N, d1, d2, ...]
135         argmax_dim: dimension to apply
136 
137     Return:
138         A tensor with categorical labels [N, d2, ...]
139 
140     Example:
141         >>> x = torch.tensor([[0.2, 0.5], [0.9, 0.1]])
142         >>> to_categorical(x)
143         tensor([1, 0])
144     """
145     return torch.argmax(x, dim=argmax_dim)
146 
147 
148 def apply_to_collection(
149     data: Any,
150     dtype: Union[type, tuple],
151     function: Callable,
152     *args: Any,
153     wrong_dtype: Optional[Union[type, tuple]] = None,
154     **kwargs: Any,
155 ) -> Any:
156     """Recursively applies a function to all elements of a certain dtype.
157 
158     Args:
159         data: the collection to apply the function to
160         dtype: the given function will be applied to all elements of this dtype
161         function: the function to apply
162         *args: positional arguments (will be forwarded to call of ``function``)
163         wrong_dtype: the given function won't be applied if this type is specified and the given collections is of
164             the :attr:`wrong_type` even if it is of type :attr`dtype`
165         **kwargs: keyword arguments (will be forwarded to call of ``function``)
166 
167     Returns:
168         the resulting collection
169 
170     Example:
171         >>> apply_to_collection(torch.tensor([8, 0, 2, 6, 7]), dtype=Tensor, function=lambda x: x ** 2)
172         tensor([64,  0,  4, 36, 49])
173         >>> apply_to_collection([8, 0, 2, 6, 7], dtype=int, function=lambda x: x ** 2)
174         [64, 0, 4, 36, 49]
175         >>> apply_to_collection(dict(abc=123), dtype=int, function=lambda x: x ** 2)
176         {'abc': 15129}
177     """
178     elem_type = type(data)
179 
180     # Breaking condition
181     if isinstance(data, dtype) and (wrong_dtype is None or not isinstance(data, wrong_dtype)):
182         return function(data, *args, **kwargs)
183 
184     # Recursively apply to collection items
185     if isinstance(data, Mapping):
186         return elem_type({k: apply_to_collection(v, dtype, function, *args, **kwargs) for k, v in data.items()})
187 
188     if isinstance(data, tuple) and hasattr(data, "_fields"):  # named tuple
189         return elem_type(*(apply_to_collection(d, dtype, function, *args, **kwargs) for d in data))
190 
191     if isinstance(data, Sequence) and not isinstance(data, str):
192         return elem_type([apply_to_collection(d, dtype, function, *args, **kwargs) for d in data])
193 
194     # data is neither of dtype, nor a collection
195     return data
196 
197 
198 def _squeeze_scalar_element_tensor(x: Tensor) -> Tensor:
199     return x.squeeze() if x.numel() == 1 else x
200 
201 
202 def _squeeze_if_scalar(data: Any) -> Any:
203     return apply_to_collection(data, Tensor, _squeeze_scalar_element_tensor)
204 
205 
206 def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:
207     """PyTorch currently does not support``torch.bincount`` for:
208 
209         - deterministic mode on GPU.
210         - MPS devices
211 
212     This implementation fallback to a for-loop counting occurrences in that case.
213 
214     Args:
215         x: tensor to count
216         minlength: minimum length to count
217 
218     Returns:
219         Number of occurrences for each unique element in x
220     """
221     if minlength is None:
222         minlength = len(torch.unique(x))
223     if torch.are_deterministic_algorithms_enabled() or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:
224         output = torch.zeros(minlength, device=x.device, dtype=torch.long)
225         for i in range(minlength):
226             output[i] = (x == i).sum()
227         return output
228     z = torch.zeros(minlength, device=x.device, dtype=x.dtype)
229     return z.index_add_(0, x, torch.ones_like(x))
230 
231 
232 def _flexible_bincount(x: Tensor) -> Tensor:
233     """Similar to `_bincount`, but works also with tensor that do not contain continuous values.
234 
235     Args:
236         x: tensor to count
237 
238     Returns:
239         Number of occurrences for each unique element in x
240     """
241 
242     # make sure elements in x start from 0
243     x = x - x.min()
244     unique_x = torch.unique(x)
245 
246     output = _bincount(x, minlength=torch.max(unique_x) + 1)
247     # remove zeros from output tensor
248     return output[unique_x]
249 
250 
251 def allclose(tensor1: Tensor, tensor2: Tensor) -> bool:
252     """Wrapper of torch.allclose that is robust towards dtype difference."""
253     if tensor1.dtype != tensor2.dtype:
254         tensor2 = tensor2.to(dtype=tensor1.dtype)
255     return torch.allclose(tensor1, tensor2)
256
```
--- 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
@@ -225,8 +225,7 @@
         for i in range(minlength):
             output[i] = (x == i).sum()
         return output
-    z = torch.zeros(minlength, device=x.device, dtype=x.dtype)
-    return z.index_add_(0, x, torch.ones_like(x))
+    return torch.bincount(x, minlength=minlength)
 
 
 def _flexible_bincount(x: Tensor) -> Tensor:

{"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@@ -225,8 +225,7 @@\n         for i in range(minlength):\n             output[i] = (x == i).sum()\n         return output\n-    z = torch.zeros(minlength, device=x.device, dtype=x.dtype)\n-    return z.index_add_(0, x, torch.ones_like(x))\n+    return torch.bincount(x, minlength=minlength)\n \n \n def _flexible_bincount(x: Tensor) -> Tensor:\n", "issue": "MatthewsCorrCoef has massive slow down on both CPU and especially GPU\n## \ud83d\udc1b Bug\r\n\r\nIt's extremely slow to compute the Matthew Correlation Coefficient since torchmetrics == 0.10.0.\r\n\r\n### To Reproduce\r\n\r\nDuring my testing, I am seeing massive slow down when computing the Matthew Correlation Coefficient, especially on the GPU (using pytorch-lightning to build, train and test a deep learning model's performance on this metric).\r\n\r\nI have compiled a code sample (next section) and below shows the result over different versions of torchmetrics. \r\n\r\nON CPU:\r\ntorchmetrics 0.7.3: 0.93961 s\r\ntorchmetrics 0.8.0: 0.93549 s\r\ntorchmetrics 0.8.2: 0.94494 s\r\ntorchmetrics 0.9.0: 0.92856 s\r\ntorchmetrics 0.9.2: 0.93682 s\r\ntorchmetrics 0.10.0: 1.10903 s\r\n\r\nON GPU:\r\ntorchmetrics 0.7.3: 0.11444 s\r\ntorchmetrics 0.8.0: 0.11682 s\r\ntorchmetrics 0.8.2: 0.11425 s\r\ntorchmetrics 0.9.0: 0.11433 s\r\ntorchmetrics 0.9.2: 0.11410 s\r\ntorchmetrics 0.10.0: 359.30208 s\r\n\r\nSo yeah testing over thousands of batches now takes almost a week for me to complete :rofl: Please take a look at this soon.\r\n\r\n#### Code sample\r\n\r\n```\r\nfrom tqdm.auto import tqdm\r\nimport time\r\nimport torch\r\nfrom torchmetrics import MatthewsCorrCoef\r\n\r\ntorch.manual_seed(1)\r\n\r\nb, h, w = 10, 1080, 1920\r\ndevice = \"cpu\"\r\n\r\ndef generate(b, h, w):\r\n    prob = torch.rand(b, h, w).to(device)\r\n    truth = torch.randint(0, 2, (b, h, w)).to(device)\r\n    return prob, truth\r\n\r\nbatches = []\r\nfor _ in range(10):\r\n    batches.append(generate(b, h, w))\r\n\r\nmcc = MatthewsCorrCoef(num_classes=2).to(device)\r\n\r\nt1 = time.time()\r\nfor detections, targets in tqdm(batches):\r\n    mcc.update(detections, targets)\r\nprint(f\"{time.time() - t1:.5f}\")\r\n```\r\n\r\n### Expected behavior\r\n\r\nIt's supposed to be faster in the range of ~0.1 second for GPU and ~0.9 second for CPU in this naive benchmark.\r\n\r\n### Environment\r\n\r\n- GPU: NVIDIA RTX 3090\r\n- TorchMetrics version (`pip`): 0.10.0\r\n- Python & PyTorch Version (e.g., 1.0): python=3.9, pytorch=1.12.1\r\n- Any other relevant information such as OS (e.g., Linux): ubuntu18.04, (same behavior with 20.04)\n", "before_files": [{"content": "# Copyright The PyTorch 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.\nfrom typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12\n\nMETRIC_EPS = 1e-6\n\n\ndef dim_zero_cat(x: Union[Tensor, List[Tensor]]) -> Tensor:\n    \"\"\"Concatenation along the zero dimension.\"\"\"\n    x = x if isinstance(x, (list, tuple)) else [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) -> Dict:\n    \"\"\"Flatten dict of dicts into single dict.\"\"\"\n    new_dict = {}\n    for key, value in x.items():\n        if isinstance(value, dict):\n            for k, v in value.items():\n                new_dict[k] = v\n        else:\n            new_dict[key] = value\n    return new_dict\n\n\ndef to_onehot(\n    label_tensor: Tensor,\n    num_classes: Optional[int] = None,\n) -> Tensor:\n    \"\"\"Converts 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    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    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    \"\"\"Converts 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    return torch.argmax(x, dim=argmax_dim)\n\n\ndef apply_to_collection(\n    data: Any,\n    dtype: Union[type, tuple],\n    function: Callable,\n    *args: Any,\n    wrong_dtype: Optional[Union[type, tuple]] = None,\n    **kwargs: Any,\n) -> Any:\n    \"\"\"Recursively applies a function to all elements of a certain dtype.\n\n    Args:\n        data: the collection to apply the function to\n        dtype: the given function will be applied to all elements of this dtype\n        function: the function to apply\n        *args: positional arguments (will be forwarded to call of ``function``)\n        wrong_dtype: the given function won't be applied if this type is specified and the given collections is of\n            the :attr:`wrong_type` even if it is of type :attr`dtype`\n        **kwargs: keyword arguments (will be forwarded to call of ``function``)\n\n    Returns:\n        the resulting collection\n\n    Example:\n        >>> apply_to_collection(torch.tensor([8, 0, 2, 6, 7]), dtype=Tensor, function=lambda x: x ** 2)\n        tensor([64,  0,  4, 36, 49])\n        >>> apply_to_collection([8, 0, 2, 6, 7], dtype=int, function=lambda x: x ** 2)\n        [64, 0, 4, 36, 49]\n        >>> apply_to_collection(dict(abc=123), dtype=int, function=lambda x: x ** 2)\n        {'abc': 15129}\n    \"\"\"\n    elem_type = type(data)\n\n    # Breaking condition\n    if isinstance(data, dtype) and (wrong_dtype is None or not isinstance(data, wrong_dtype)):\n        return function(data, *args, **kwargs)\n\n    # Recursively apply to collection items\n    if isinstance(data, Mapping):\n        return elem_type({k: apply_to_collection(v, dtype, function, *args, **kwargs) for k, v in data.items()})\n\n    if isinstance(data, tuple) and hasattr(data, \"_fields\"):  # named tuple\n        return elem_type(*(apply_to_collection(d, dtype, function, *args, **kwargs) for d in data))\n\n    if isinstance(data, Sequence) and not isinstance(data, str):\n        return elem_type([apply_to_collection(d, dtype, function, *args, **kwargs) for d in data])\n\n    # data is neither of dtype, nor a collection\n    return data\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    \"\"\"PyTorch currently does not support``torch.bincount`` for:\n\n        - deterministic mode on GPU.\n        - MPS devices\n\n    This implementation fallback to a for-loop counting occurrences in that case.\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    if minlength is None:\n        minlength = len(torch.unique(x))\n    if torch.are_deterministic_algorithms_enabled() or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:\n        output = torch.zeros(minlength, device=x.device, dtype=torch.long)\n        for i in range(minlength):\n            output[i] = (x == i).sum()\n        return output\n    z = torch.zeros(minlength, device=x.device, dtype=x.dtype)\n    return z.index_add_(0, x, torch.ones_like(x))\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)\n    # remove zeros from output tensor\n    return output[unique_x]\n\n\ndef allclose(tensor1: Tensor, tensor2: Tensor) -> bool:\n    \"\"\"Wrapper of torch.allclose that is 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 PyTorch 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.\nfrom typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12\n\nMETRIC_EPS = 1e-6\n\n\ndef dim_zero_cat(x: Union[Tensor, List[Tensor]]) -> Tensor:\n    \"\"\"Concatenation along the zero dimension.\"\"\"\n    x = x if isinstance(x, (list, tuple)) else [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) -> Dict:\n    \"\"\"Flatten dict of dicts into single dict.\"\"\"\n    new_dict = {}\n    for key, value in x.items():\n        if isinstance(value, dict):\n            for k, v in value.items():\n                new_dict[k] = v\n        else:\n            new_dict[key] = value\n    return new_dict\n\n\ndef to_onehot(\n    label_tensor: Tensor,\n    num_classes: Optional[int] = None,\n) -> Tensor:\n    \"\"\"Converts 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    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    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    \"\"\"Converts 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    return torch.argmax(x, dim=argmax_dim)\n\n\ndef apply_to_collection(\n    data: Any,\n    dtype: Union[type, tuple],\n    function: Callable,\n    *args: Any,\n    wrong_dtype: Optional[Union[type, tuple]] = None,\n    **kwargs: Any,\n) -> Any:\n    \"\"\"Recursively applies a function to all elements of a certain dtype.\n\n    Args:\n        data: the collection to apply the function to\n        dtype: the given function will be applied to all elements of this dtype\n        function: the function to apply\n        *args: positional arguments (will be forwarded to call of ``function``)\n        wrong_dtype: the given function won't be applied if this type is specified and the given collections is of\n            the :attr:`wrong_type` even if it is of type :attr`dtype`\n        **kwargs: keyword arguments (will be forwarded to call of ``function``)\n\n    Returns:\n        the resulting collection\n\n    Example:\n        >>> apply_to_collection(torch.tensor([8, 0, 2, 6, 7]), dtype=Tensor, function=lambda x: x ** 2)\n        tensor([64,  0,  4, 36, 49])\n        >>> apply_to_collection([8, 0, 2, 6, 7], dtype=int, function=lambda x: x ** 2)\n        [64, 0, 4, 36, 49]\n        >>> apply_to_collection(dict(abc=123), dtype=int, function=lambda x: x ** 2)\n        {'abc': 15129}\n    \"\"\"\n    elem_type = type(data)\n\n    # Breaking condition\n    if isinstance(data, dtype) and (wrong_dtype is None or not isinstance(data, wrong_dtype)):\n        return function(data, *args, **kwargs)\n\n    # Recursively apply to collection items\n    if isinstance(data, Mapping):\n        return elem_type({k: apply_to_collection(v, dtype, function, *args, **kwargs) for k, v in data.items()})\n\n    if isinstance(data, tuple) and hasattr(data, \"_fields\"):  # named tuple\n        return elem_type(*(apply_to_collection(d, dtype, function, *args, **kwargs) for d in data))\n\n    if isinstance(data, Sequence) and not isinstance(data, str):\n        return elem_type([apply_to_collection(d, dtype, function, *args, **kwargs) for d in data])\n\n    # data is neither of dtype, nor a collection\n    return data\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    \"\"\"PyTorch currently does not support``torch.bincount`` for:\n\n        - deterministic mode on GPU.\n        - MPS devices\n\n    This implementation fallback to a for-loop counting occurrences in that case.\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    if minlength is None:\n        minlength = len(torch.unique(x))\n    if torch.are_deterministic_algorithms_enabled() or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:\n        output = torch.zeros(minlength, device=x.device, dtype=torch.long)\n        for i in range(minlength):\n            output[i] = (x == i).sum()\n        return output\n    return torch.bincount(x, minlength=minlength)\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)\n    # remove zeros from output tensor\n    return output[unique_x]\n\n\ndef allclose(tensor1: Tensor, tensor2: Tensor) -> bool:\n    \"\"\"Wrapper of torch.allclose that is 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"}]}