Repository: Lightning-AI/torchmetrics

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


Discussion:
Hi! thanks for your contribution!, great first issue!
Hi @mfatih7 ,
Thanks for the issue. Would you be interested in doing so or at least give us a hint on what's not working on TPU?
Hi @justusschock

I can give information about the situation in the TPU.
I have a PyTorch project in which I use different kinds of models for classification.
In the program flow, I import your module as

`from torchmetrics import ConfusionMatrix`

I instantiate an instance using

`confmat = ConfusionMatrix(task="binary", num_classes=2).to(device)`

When the device is GPU it works without any problem.
When the device is TPU confmat causes [XLA compilations](https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md)

Actually, it is a normal situation.
PyTorch XLA team welcomes lowering requests for torch functions.
The functions causing compilations are `_unique2` and `bincount`.
You can check my issue in [XLA issue](https://github.com/pytorch/xla/issues/4511).

Maybe you need to implement your module without these functions.

I can give more information with pleasure.






To prevent XLA compilations temporarily, I am using the simple function below.

```
def get_confusion_matrix_for_xla(outputs, labels):
    
    confusion_matrix = torch.zeros( 2, 2, dtype=torch.int64, device = params.DEVICE)
    
    confusion_matrix[0,0] = torch.sum( torch.logical_and( labels < 0.5, outputs < 0.5) )
    confusion_matrix[0,1] = torch.sum( torch.logical_and( labels < 0.5, outputs >= 0.5) )
    confusion_matrix[1,0] = torch.sum( torch.logical_and( labels >= 0.5, outputs < 0.5) )
    confusion_matrix[1,1] = torch.sum( torch.logical_and( labels >= 0.5, outputs >= 0.5) )
    
    return confusion_matrix
```

Hi @mfatih7,
Is it possible to reproduce this behaviour in a colab notebook?
While we cannot rewrite our hole framework to support XLA, we could probably implement fallback solutions on specific devices. For example `_bincount` is a function that is not support by MPS accelerator yet, and therefore we have a fallback solution for that:
https://github.com/Lightning-AI/metrics/blob/5d4ffe01aa09b7108f7e0e4034748bdfd64bf5f9/src/torchmetrics/utilities/data.py#L206-L228
hi @SkafteNicki 

I could not get the warnings with code purely written in the COLAB notebook.
But [here](https://drive.google.com/drive/folders/1SuOHBUyV3ThMCyl25kVvEO__jNl61JJP?usp=sharing) is a .py file and a COLAB notebook.
You can easily see the warnings.

I can give more support if needed.
More debugging options are available.

Do not forget to select TPU from COLAB settings

Resolving patch:
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 @@ 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()
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")


Test patch: