Repository: Lightning-AI/torchmetrics

Issue:
JaccardIndex mps fails also with PYTORCH_ENABLE_MPS_FALLBACK=1
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

JaccardIndex call fails with M1 mps backend since "aten::bincount" is not currently supported. This happens also when PYTORCH_ENABLE_MPS_FALLBACK=1 is set environment variable.

### To Reproduce

```
IoU = JaccardIndex(num_classes=2, ignore_index=0)`
IoU = IoU.to("mps")`
for batch in train_dl:
    x, mask = batch
    x = x.to("mps")
    mask = mask.to("mps")
    logits = model(x)
    preds = torch.sigmoid(logits)
    iou = IoU(preds, mask)
```
   
UserWarning: The operator 'aten::bincount' is not currently supported on the MPS backend and will fall back to run on the CPU. This may have performance implications. (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/mps/MPSFallback.mm:11.)
  return torch.bincount(x, minlength=minlength)


### Environment

- TorchMetrics version (`pip`): 0.9.3
- Python & PyTorch Version (e.g., 1.0): 1.12.1
- Any other relevant information such as OS (e.g., Linux): macOS with M1



Discussion:
Hi! thanks for your contribution!, great first issue!
@justusschock should we do anything to support this explicitly or do we need to wait for support from PyTorch side? 
Hi @SkafteNicki, I think as a temporary fix, we may adjust the following part of code:
```diff
def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:
    """``torch.bincount`` currently does not support deterministic mode on GPU.
    This implementation fallback to a for-loop counting occurrences in that case.
    Args:
        x: tensor to count
        minlength: minimum length to count
    Returns:
        Number of occurrences for each unique element in x
    """
-    if x.is_cuda and deterministic():
+   if (x.is_cuda and deterministic()) or x.is_mps:
        if minlength is None:
            minlength = len(torch.unique(x))
        output = torch.zeros(minlength, device=x.device, dtype=torch.long)
        for i in range(minlength):
            output[i] = (x == i).sum()
        return output
    else:
        return torch.bincount(x, minlength=minlength)
```
Also, there are some other operations which are not yet supported on `MPS`. I can send a short PR fixing those issues and open an issue to check those handlings if they're still needed once `torch` is updated.
@stancld feel free to do so :)
@Borda Do you think it'd be feasible and wise to run our test suites also on an M1 runner?
https://github.blog/changelog/2022-08-09-github-actions-self-hosted-runners-now-support-apple-m1-hardware/

Resolving patch:
diff --git a/src/torchmetrics/functional/classification/jaccard.py b/src/torchmetrics/functional/classification/jaccard.py
--- a/src/torchmetrics/functional/classification/jaccard.py
+++ b/src/torchmetrics/functional/classification/jaccard.py
@@ -62,7 +62,7 @@ def _jaccard_from_confmat(
 
         # If this class is absent in both target AND pred (union == 0), then use the absent_score for this class.
         scores = intersection.float() / union.float()
-        scores[union == 0] = absent_score
+        scores = scores.where(union != 0, torch.tensor(absent_score, dtype=scores.dtype, device=scores.device))
 
         if ignore_index is not None and 0 <= ignore_index < num_classes:
             scores = torch.cat(
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,12 @@
 import torch
 from torch import Tensor, tensor
 
-from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_6, _TORCH_GREATER_EQUAL_1_7, _TORCH_GREATER_EQUAL_1_8
+from torchmetrics.utilities.imports import (
+    _TORCH_GREATER_EQUAL_1_6,
+    _TORCH_GREATER_EQUAL_1_7,
+    _TORCH_GREATER_EQUAL_1_8,
+    _TORCH_GREATER_EQUAL_1_12,
+)
 
 if _TORCH_GREATER_EQUAL_1_8:
     deterministic = torch.are_deterministic_algorithms_enabled
@@ -242,7 +247,10 @@ def _squeeze_if_scalar(data: Any) -> Any:
 
 
 def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:
-    """``torch.bincount`` currently does not support deterministic mode on GPU.
+    """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.
 
@@ -253,7 +261,7 @@ def _bincount(x: Tensor, minlength: Optional[int] = None) -> Tensor:
     Returns:
         Number of occurrences for each unique element in x
     """
-    if x.is_cuda and deterministic():
+    if x.is_cuda and deterministic() or _TORCH_GREATER_EQUAL_1_12 and x.is_mps:
         if minlength is None:
             minlength = len(torch.unique(x))
         output = torch.zeros(minlength, device=x.device, dtype=torch.long)
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
@@ -104,6 +104,7 @@ def _compare_version(package: str, op: Callable, version: str) -> Optional[bool]
 _TORCH_GREATER_EQUAL_1_8: Optional[bool] = _compare_version("torch", operator.ge, "1.8.0")
 _TORCH_GREATER_EQUAL_1_10: Optional[bool] = _compare_version("torch", operator.ge, "1.10.0")
 _TORCH_GREATER_EQUAL_1_11: Optional[bool] = _compare_version("torch", operator.ge, "1.11.0")
+_TORCH_GREATER_EQUAL_1_12: Optional[bool] = _compare_version("torch", operator.ge, "1.12.0")
 
 _JIWER_AVAILABLE: bool = _package_available("jiwer")
 _NLTK_AVAILABLE: bool = _package_available("nltk")


Test patch: