We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
XLA support
Hello

Up to now, I was using `torchmetrics` in my training scripts running on GPUs.
Now I want to use Google Tensor Processing Units in my work.
For the last few days, I am observing that `torchmetrics` is not compatible with XLA library.
`torchmetrics` needs to be lowered for TPU support.

best regards

--- 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     return torch.bincount(x, minlength=minlength)
229 
230 
231 def _flexible_bincount(x: Tensor) -> Tensor:
232     """Similar to `_bincount`, but works also with tensor that do not contain continuous values.
233 
234     Args:
235         x: tensor to count
236 
237     Returns:
238         Number of occurrences for each unique element in x
239     """
240 
241     # make sure elements in x start from 0
242     x = x - x.min()
243     unique_x = torch.unique(x)
244 
245     output = _bincount(x, minlength=torch.max(unique_x) + 1)
246     # remove zeros from output tensor
247     return output[unique_x]
248 
249 
250 def allclose(tensor1: Tensor, tensor2: Tensor) -> bool:
251     """Wrapper of torch.allclose that is robust towards dtype difference."""
252     if tensor1.dtype != tensor2.dtype:
253         tensor2 = tensor2.to(dtype=tensor1.dtype)
254     return torch.allclose(tensor1, tensor2)
255
```

Path: `src/torchmetrics/utilities/imports.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 """Import utilities."""
15 import operator
16 from typing import Optional
17 
18 from lightning_utilities.core.imports import compare_version, package_available
19 
20 _TORCH_LOWER_1_12_DEV: Optional[bool] = compare_version("torch", operator.lt, "1.12.0.dev")
21 _TORCH_GREATER_EQUAL_1_9: Optional[bool] = compare_version("torch", operator.ge, "1.9.0")
22 _TORCH_GREATER_EQUAL_1_10: Optional[bool] = compare_version("torch", operator.ge, "1.10.0")
23 _TORCH_GREATER_EQUAL_1_11: Optional[bool] = compare_version("torch", operator.ge, "1.11.0")
24 _TORCH_GREATER_EQUAL_1_12: Optional[bool] = compare_version("torch", operator.ge, "1.12.0")
25 
26 _JIWER_AVAILABLE: bool = package_available("jiwer")
27 _NLTK_AVAILABLE: bool = package_available("nltk")
28 _ROUGE_SCORE_AVAILABLE: bool = package_available("rouge_score")
29 _BERTSCORE_AVAILABLE: bool = package_available("bert_score")
30 _SCIPY_AVAILABLE: bool = package_available("scipy")
31 _TORCH_FIDELITY_AVAILABLE: bool = package_available("torch_fidelity")
32 _LPIPS_AVAILABLE: bool = package_available("lpips")
33 _PYCOCOTOOLS_AVAILABLE: bool = package_available("pycocotools")
34 _TORCHVISION_AVAILABLE: bool = package_available("torchvision")
35 _TORCHVISION_GREATER_EQUAL_0_8: Optional[bool] = compare_version("torchvision", operator.ge, "0.8.0")
36 _TQDM_AVAILABLE: bool = package_available("tqdm")
37 _TRANSFORMERS_AVAILABLE: bool = package_available("transformers")
38 _PESQ_AVAILABLE: bool = package_available("pesq")
39 _SACREBLEU_AVAILABLE: bool = package_available("sacrebleu")
40 _REGEX_AVAILABLE: bool = package_available("regex")
41 _PYSTOI_AVAILABLE: bool = package_available("pystoi")
42 _FAST_BSS_EVAL_AVAILABLE: bool = package_available("fast_bss_eval")
43 _MATPLOTLIB_AVAILABLE: bool = package_available("matplotlib")
44 _MULTIPROCESSING_AVAILABLE: bool = package_available("multiprocessing")
45
```
--- 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
@@ -16,7 +16,7 @@
 import torch
 from torch import Tensor
 
-from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12
+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12, _XLA_AVAILABLE
 
 METRIC_EPS = 1e-6
 
@@ -220,7 +220,7 @@
     """
     if minlength is None:
         minlength = len(torch.unique(x))
-    if torch.are_deterministic_algorithms_enabled() or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:
+    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()
diff --git a/src/torchmetrics/utilities/imports.py b/src/torchmetrics/utilities/imports.py
--- a/src/torchmetrics/utilities/imports.py
+++ b/src/torchmetrics/utilities/imports.py
@@ -42,3 +42,4 @@
 _FAST_BSS_EVAL_AVAILABLE: bool = package_available("fast_bss_eval")
 _MATPLOTLIB_AVAILABLE: bool = package_available("matplotlib")
 _MULTIPROCESSING_AVAILABLE: bool = package_available("multiprocessing")
+_XLA_AVAILABLE: bool = package_available("torch_xla")

{"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@@ -16,7 +16,7 @@\n import torch\n from torch import Tensor\n \n-from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12\n+from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_12, _XLA_AVAILABLE\n \n METRIC_EPS = 1e-6\n \n@@ -220,7 +220,7 @@\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+    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()\ndiff --git a/src/torchmetrics/utilities/imports.py b/src/torchmetrics/utilities/imports.py\n--- a/src/torchmetrics/utilities/imports.py\n+++ b/src/torchmetrics/utilities/imports.py\n@@ -42,3 +42,4 @@\n _FAST_BSS_EVAL_AVAILABLE: bool = package_available(\"fast_bss_eval\")\n _MATPLOTLIB_AVAILABLE: bool = package_available(\"matplotlib\")\n _MULTIPROCESSING_AVAILABLE: bool = package_available(\"multiprocessing\")\n+_XLA_AVAILABLE: bool = package_available(\"torch_xla\")\n", "issue": "XLA support\nHello\r\n\r\nUp to now, I was using `torchmetrics` in my training scripts running on GPUs.\r\nNow I want to use Google Tensor Processing Units in my work.\r\nFor the last few days, I am observing that `torchmetrics` is not compatible with XLA library.\r\n`torchmetrics` needs to be lowered for TPU support.\r\n\r\nbest regards\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    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"}, {"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.\n\"\"\"Import utilities.\"\"\"\nimport operator\nfrom typing import Optional\n\nfrom lightning_utilities.core.imports import compare_version, package_available\n\n_TORCH_LOWER_1_12_DEV: Optional[bool] = compare_version(\"torch\", operator.lt, \"1.12.0.dev\")\n_TORCH_GREATER_EQUAL_1_9: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.9.0\")\n_TORCH_GREATER_EQUAL_1_10: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.10.0\")\n_TORCH_GREATER_EQUAL_1_11: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.11.0\")\n_TORCH_GREATER_EQUAL_1_12: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.12.0\")\n\n_JIWER_AVAILABLE: bool = package_available(\"jiwer\")\n_NLTK_AVAILABLE: bool = package_available(\"nltk\")\n_ROUGE_SCORE_AVAILABLE: bool = package_available(\"rouge_score\")\n_BERTSCORE_AVAILABLE: bool = package_available(\"bert_score\")\n_SCIPY_AVAILABLE: bool = package_available(\"scipy\")\n_TORCH_FIDELITY_AVAILABLE: bool = package_available(\"torch_fidelity\")\n_LPIPS_AVAILABLE: bool = package_available(\"lpips\")\n_PYCOCOTOOLS_AVAILABLE: bool = package_available(\"pycocotools\")\n_TORCHVISION_AVAILABLE: bool = package_available(\"torchvision\")\n_TORCHVISION_GREATER_EQUAL_0_8: Optional[bool] = compare_version(\"torchvision\", operator.ge, \"0.8.0\")\n_TQDM_AVAILABLE: bool = package_available(\"tqdm\")\n_TRANSFORMERS_AVAILABLE: bool = package_available(\"transformers\")\n_PESQ_AVAILABLE: bool = package_available(\"pesq\")\n_SACREBLEU_AVAILABLE: bool = package_available(\"sacrebleu\")\n_REGEX_AVAILABLE: bool = package_available(\"regex\")\n_PYSTOI_AVAILABLE: bool = package_available(\"pystoi\")\n_FAST_BSS_EVAL_AVAILABLE: bool = package_available(\"fast_bss_eval\")\n_MATPLOTLIB_AVAILABLE: bool = package_available(\"matplotlib\")\n_MULTIPROCESSING_AVAILABLE: bool = package_available(\"multiprocessing\")\n", "path": "src/torchmetrics/utilities/imports.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, _XLA_AVAILABLE\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 _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 _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"}, {"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.\n\"\"\"Import utilities.\"\"\"\nimport operator\nfrom typing import Optional\n\nfrom lightning_utilities.core.imports import compare_version, package_available\n\n_TORCH_LOWER_1_12_DEV: Optional[bool] = compare_version(\"torch\", operator.lt, \"1.12.0.dev\")\n_TORCH_GREATER_EQUAL_1_9: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.9.0\")\n_TORCH_GREATER_EQUAL_1_10: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.10.0\")\n_TORCH_GREATER_EQUAL_1_11: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.11.0\")\n_TORCH_GREATER_EQUAL_1_12: Optional[bool] = compare_version(\"torch\", operator.ge, \"1.12.0\")\n\n_JIWER_AVAILABLE: bool = package_available(\"jiwer\")\n_NLTK_AVAILABLE: bool = package_available(\"nltk\")\n_ROUGE_SCORE_AVAILABLE: bool = package_available(\"rouge_score\")\n_BERTSCORE_AVAILABLE: bool = package_available(\"bert_score\")\n_SCIPY_AVAILABLE: bool = package_available(\"scipy\")\n_TORCH_FIDELITY_AVAILABLE: bool = package_available(\"torch_fidelity\")\n_LPIPS_AVAILABLE: bool = package_available(\"lpips\")\n_PYCOCOTOOLS_AVAILABLE: bool = package_available(\"pycocotools\")\n_TORCHVISION_AVAILABLE: bool = package_available(\"torchvision\")\n_TORCHVISION_GREATER_EQUAL_0_8: Optional[bool] = compare_version(\"torchvision\", operator.ge, \"0.8.0\")\n_TQDM_AVAILABLE: bool = package_available(\"tqdm\")\n_TRANSFORMERS_AVAILABLE: bool = package_available(\"transformers\")\n_PESQ_AVAILABLE: bool = package_available(\"pesq\")\n_SACREBLEU_AVAILABLE: bool = package_available(\"sacrebleu\")\n_REGEX_AVAILABLE: bool = package_available(\"regex\")\n_PYSTOI_AVAILABLE: bool = package_available(\"pystoi\")\n_FAST_BSS_EVAL_AVAILABLE: bool = package_available(\"fast_bss_eval\")\n_MATPLOTLIB_AVAILABLE: bool = package_available(\"matplotlib\")\n_MULTIPROCESSING_AVAILABLE: bool = package_available(\"multiprocessing\")\n_XLA_AVAILABLE: bool = package_available(\"torch_xla\")\n", "path": "src/torchmetrics/utilities/imports.py"}]}