We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Validation with torchmetrics extremely slow
### Bug description

Hi all,

I recently tried to implement a DeepLabV3 training pipeline. I wanted to use the build-in `torchmetrics.JaccardIndex` as my evaluation metric. My `LightningModule`  looks like this:

```python
import torchmetrics 
from pytorch_lightning import LightningModule
from torchvision.models.segmentation.deeplabv3 import deeplabv3_resnet50


class DeepLabV3LightningModule(LightningModule):
    def __init__(self):
        self.model = deeplabv3_resnet50(
            num_classes=38,
            aux_loss=False
        )
        self.loss = nn.CrossEntropyLoss(ignore_index=255, reduction="mean")
        self.iou_metric = torchmetrics.JaccardIndex(
            task="multiclass", 
            threshold=0.5, 
            num_classes=38,
            average="macro",
        )

    def training_step(self, batch, batch_idx):
        imgs, masks = batch
        out = self.model(imgs)
        preds = out["out"]
        loss = self.loss(preds, masks)       
        return loss

    def validation_step(self, batch, batch_idx):
        imgs, masks = batch
        out = self.model(imgs)
        preds = out["out"]
        loss = self.loss(preds, masks)
        preds = torch.softmax(preds, dim=1)
        pred_labels = torch.argmax(preds, dim=1)
        
        # measure runtime of metric update
        start = timer()
        self.iou_metric.update(pred_labels, masks)
        elapsed = timer() - start
        return elapsed

    def validation_epoch_end(self, outputs):
        avg_runtime = round(mean(outputs), 4)
        print(f"GPU {self.local_rank}: {avg_runtime} seconds")

```

When using this validation procedure, it is extremely slow. On average, the update step of the metric takes 23.4 seconds. However, the first 3 updates are very fast (<1 second), then they become slow.

I tried to reproduce this behavior in a MWE:

```python
from timeit import default_timer as timer
from statistics import mean
import torchmetrics
import torch

num_classes = 38

iou_metric = torchmetrics.JaccardIndex(
    task="multiclass",
    threshold=0.5, 
    num_classes=num_classes,
    average="macro"
).to("cuda")

# dummy labels in shape [b, h, w]
label_mask = torch.randint(low=0, high=num_classes-1, size=(8, 480, 640), device="cuda")

# dummy predicted labels in shape [b, h, w]
pred_mask = torch.randint(low=0, high=num_classes-1, size=(8, 480, 640), device="cuda")


runtime_hist = []
for i in range(100):
    start = timer()
    iou_metric.update(label_mask, pred_mask)
    elapsed = timer() - start
    runtime_hist.append(elapsed)


avg_runtime = round(mean(runtime_hist), 2)
print(avg_runtime)
```

Here I get an average update duration of 0.03 seconds, so I do not encounter the extremely slow update as in my `LightningModule` above. To me this looks like there is something wrong. At this point, I am not sure if thi


Here some training information for my pytorch-lightning training pipeline:
- OS: Ubuntu 20.04.4
- CUDA 11.3
- DDP training strategy
- GPUs: 4x V100
- batch size: 8
- image size (width x height): 640 x 480
- number of workers in dataloader: 8

My package versions:
- pytorch lightning: 1.8.4.post0 (installed via pip)
- torch: 1.13.0
- torchvision: 0.14.0
- torchmetrics: 0.11.0
- numpy: 11.23.5

Thanks so much!
Lukas

### How to reproduce the bug

_No response_

### Error messages and logs

_No response_

### Environment

_No response_

### More info

_No response_

cc @carmocca @justusschock @awaelchli @borda

--- 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`` for:
173 
174         - deterministic mode on GPU.
175         - MPS devices
176 
177     This implementation fallback to a for-loop counting occurrences in that case.
178 
179     Args:
180         x: tensor to count
181         minlength: minimum length to count
182 
183     Returns:
184         Number of occurrences for each unique element in x
185 
186     Example:
187         >>> x = torch.tensor([0,0,0,1,1,2,2,2,2])
188         >>> _bincount(x, minlength=3)
189         tensor([3, 2, 4])
190 
191     """
192     if minlength is None:
193         minlength = len(torch.unique(x))
194     if torch.are_deterministic_algorithms_enabled() or _XLA_AVAILABLE or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:
195         output = torch.zeros(minlength, device=x.device, dtype=torch.long)
196         for i in range(minlength):
197             output[i] = (x == i).sum()
198         return output
199     return torch.bincount(x, minlength=minlength)
200 
201 
202 def _cumsum(x: Tensor, dim: Optional[int] = 0, dtype: Optional[torch.dtype] = None) -> Tensor:
203     if torch.are_deterministic_algorithms_enabled() and x.is_cuda and x.is_floating_point() and sys.platform != "win32":
204         rank_zero_warn(
205             "You are trying to use a metric in deterministic mode on GPU that uses `torch.cumsum`, which is currently "
206             "not supported. The tensor will be copied to the CPU memory to compute it and then copied back to GPU. "
207             "Expect some slowdowns.",
208             TorchMetricsUserWarning,
209         )
210         return x.cpu().cumsum(dim=dim, dtype=dtype).cuda()
211     return torch.cumsum(x, dim=dim, dtype=dtype)
212 
213 
214 def _flexible_bincount(x: Tensor) -> Tensor:
215     """Similar to `_bincount`, but works also with tensor that do not contain continuous values.
216 
217     Args:
218         x: tensor to count
219 
220     Returns:
221         Number of occurrences for each unique element in x
222 
223     """
224     # make sure elements in x start from 0
225     x = x - x.min()
226     unique_x = torch.unique(x)
227 
228     output = _bincount(x, minlength=torch.max(unique_x) + 1)  # type: ignore[arg-type]
229     # remove zeros from output tensor
230     return output[unique_x]
231 
232 
233 def allclose(tensor1: Tensor, tensor2: Tensor) -> bool:
234     """Wrap torch.allclose to be robust towards dtype difference."""
235     if tensor1.dtype != tensor2.dtype:
236         tensor2 = tensor2.to(dtype=tensor1.dtype)
237     return torch.allclose(tensor1, tensor2)
238
```
--- 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
@@ -169,12 +169,10 @@
 def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:
     """Implement custom bincount.
 
-    PyTorch currently does not support ``torch.bincount`` for:
-
-        - deterministic mode on GPU.
-        - MPS devices
-
-    This implementation fallback to a for-loop counting occurrences in that case.
+    PyTorch currently does not support ``torch.bincount`` when running in deterministic mode on GPU or when running
+    MPS devices or when running on XLA device. This implementation therefore falls back to using a combination of
+    `torch.arange` and `torch.eq` in these scenarios. A small performance hit can expected and higher memory consumption
+    as `[batch_size, mincount]` tensor needs to be initialized compared to native ``torch.bincount``.
 
     Args:
         x: tensor to count
@@ -191,11 +189,11 @@
     """
     if minlength is None:
         minlength = len(torch.unique(x))
+
     if torch.are_deterministic_algorithms_enabled() or _XLA_AVAILABLE or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:
-        output = torch.zeros(minlength, device=x.device, dtype=torch.long)
-        for i in range(minlength):
-            output[i] = (x == i).sum()
-        return output
+        mesh = torch.arange(minlength, device=x.device).repeat(len(x), 1)
+        return torch.eq(x.reshape(-1, 1), mesh).sum(dim=0)
+
     return torch.bincount(x, minlength=minlength)

{"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@@ -169,12 +169,10 @@\n def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:\n     \"\"\"Implement custom bincount.\n \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+    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@@ -191,11 +189,11 @@\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-        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+        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", "issue": "Validation with torchmetrics extremely slow\n### Bug description\r\n\r\nHi all,\r\n\r\nI recently tried to implement a DeepLabV3 training pipeline. I wanted to use the build-in `torchmetrics.JaccardIndex` as my evaluation metric. My `LightningModule`  looks like this:\r\n\r\n```python\r\nimport torchmetrics \r\nfrom pytorch_lightning import LightningModule\r\nfrom torchvision.models.segmentation.deeplabv3 import deeplabv3_resnet50\r\n\r\n\r\nclass DeepLabV3LightningModule(LightningModule):\r\n    def __init__(self):\r\n        self.model = deeplabv3_resnet50(\r\n            num_classes=38,\r\n            aux_loss=False\r\n        )\r\n        self.loss = nn.CrossEntropyLoss(ignore_index=255, reduction=\"mean\")\r\n        self.iou_metric = torchmetrics.JaccardIndex(\r\n            task=\"multiclass\", \r\n            threshold=0.5, \r\n            num_classes=38,\r\n            average=\"macro\",\r\n        )\r\n\r\n    def training_step(self, batch, batch_idx):\r\n        imgs, masks = batch\r\n        out = self.model(imgs)\r\n        preds = out[\"out\"]\r\n        loss = self.loss(preds, masks)       \r\n        return loss\r\n\r\n    def validation_step(self, batch, batch_idx):\r\n        imgs, masks = batch\r\n        out = self.model(imgs)\r\n        preds = out[\"out\"]\r\n        loss = self.loss(preds, masks)\r\n        preds = torch.softmax(preds, dim=1)\r\n        pred_labels = torch.argmax(preds, dim=1)\r\n        \r\n        # measure runtime of metric update\r\n        start = timer()\r\n        self.iou_metric.update(pred_labels, masks)\r\n        elapsed = timer() - start\r\n        return elapsed\r\n\r\n    def validation_epoch_end(self, outputs):\r\n        avg_runtime = round(mean(outputs), 4)\r\n        print(f\"GPU {self.local_rank}: {avg_runtime} seconds\")\r\n\r\n```\r\n\r\nWhen using this validation procedure, it is extremely slow. On average, the update step of the metric takes 23.4 seconds. However, the first 3 updates are very fast (<1 second), then they become slow.\r\n\r\nI tried to reproduce this behavior in a MWE:\r\n\r\n```python\r\nfrom timeit import default_timer as timer\r\nfrom statistics import mean\r\nimport torchmetrics\r\nimport torch\r\n\r\nnum_classes = 38\r\n\r\niou_metric = torchmetrics.JaccardIndex(\r\n    task=\"multiclass\",\r\n    threshold=0.5, \r\n    num_classes=num_classes,\r\n    average=\"macro\"\r\n).to(\"cuda\")\r\n\r\n# dummy labels in shape [b, h, w]\r\nlabel_mask = torch.randint(low=0, high=num_classes-1, size=(8, 480, 640), device=\"cuda\")\r\n\r\n# dummy predicted labels in shape [b, h, w]\r\npred_mask = torch.randint(low=0, high=num_classes-1, size=(8, 480, 640), device=\"cuda\")\r\n\r\n\r\nruntime_hist = []\r\nfor i in range(100):\r\n    start = timer()\r\n    iou_metric.update(label_mask, pred_mask)\r\n    elapsed = timer() - start\r\n    runtime_hist.append(elapsed)\r\n\r\n\r\navg_runtime = round(mean(runtime_hist), 2)\r\nprint(avg_runtime)\r\n```\r\n\r\nHere I get an average update duration of 0.03 seconds, so I do not encounter the extremely slow update as in my `LightningModule` above. To me this looks like there is something wrong. At this point, I am not sure if thi\r\n\r\n\r\nHere some training information for my pytorch-lightning training pipeline:\r\n- OS: Ubuntu 20.04.4\r\n- CUDA 11.3\r\n- DDP training strategy\r\n- GPUs: 4x V100\r\n- batch size: 8\r\n- image size (width x height): 640 x 480\r\n- number of workers in dataloader: 8\r\n\r\nMy package versions:\r\n- pytorch lightning: 1.8.4.post0 (installed via pip)\r\n- torch: 1.13.0\r\n- torchvision: 0.14.0\r\n- torchmetrics: 0.11.0\r\n- numpy: 11.23.5\r\n\r\nThanks so much!\r\nLukas\r\n\r\n### How to reproduce the bug\r\n\r\n_No response_\r\n\r\n### Error messages and logs\r\n\r\n_No response_\r\n\r\n### Environment\r\n\r\n_No response_\r\n\r\n### More info\r\n\r\n_No response_\n\ncc @carmocca @justusschock @awaelchli @borda\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`` 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    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    if torch.are_deterministic_algorithms_enabled() or _XLA_AVAILABLE 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 _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 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"}]}