Fixed bincount on XLA (#1471)

CHANGELOG.md

@@ -40,6 +40,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Fix compatibility between XLA in `_bincount` function ([#1471](https://github.com/Lightning-AI/metrics/pull/1471))
+
+
 -
 
 

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 @@ def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:
     """
     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()

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")