Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fixed bincount on XLA #1471

Merged
merged 5 commits into from Jan 31, 2023
Merged

Fixed bincount on XLA #1471

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
90d9bb3
try
SkafteNicki Jan 31, 2023
ca66422
changelog
SkafteNicki Jan 31, 2023
0cac211
Merge branch 'master' into xla_test
Borda Jan 31, 2023
37b0ae5
Merge branch 'master' into xla_test
Borda Jan 31, 2023
0d89765
chlog
Borda Jan 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/torchmetrics/utilities/data.py
View file
Open in desktop
Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down
1 change: 1 addition & 0 deletions src/torchmetrics/utilities/imports.py
View file
Open in desktop
Expand Up @@ -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")