Skip to content

Commit 61f9594

Browse files
authored
Merge branch 'main' into detect-numerical-code
2 parents 1e5c7d6 + 9c04ea6 commit 61f9594

2 files changed

Lines changed: 47 additions & 0 deletions

File tree

codeflash/verification/comparator.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,6 +296,9 @@ def comparator(orig: Any, new: Any, superset_obj=False) -> bool: # noqa: ANN001
296296
if isinstance(orig, torch.dtype):
297297
return orig == new
298298

299+
if isinstance(orig, torch.device):
300+
return orig == new
301+
299302
if HAS_PYRSISTENT:
300303
import pyrsistent # type: ignore # noqa: PGH003
301304

tests/test_comparator.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -824,6 +824,50 @@ def test_torch():
824824
assert not comparator(gg, ii)
825825

826826

827+
def test_torch_device():
828+
try:
829+
import torch # type: ignore
830+
except ImportError:
831+
pytest.skip()
832+
833+
# Test torch.device comparisons - same device type
834+
a = torch.device("cpu")
835+
b = torch.device("cpu")
836+
assert comparator(a, b)
837+
838+
# Test different device types
839+
c = torch.device("cpu")
840+
d = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
841+
if torch.cuda.is_available():
842+
assert not comparator(c, d)
843+
844+
# Test device with index
845+
e = torch.device("cpu")
846+
f = torch.device("cpu")
847+
assert comparator(e, f)
848+
849+
# Test cuda devices with different indices (if multiple GPUs available)
850+
if torch.cuda.is_available() and torch.cuda.device_count() > 1:
851+
g = torch.device("cuda:0")
852+
h = torch.device("cuda:0")
853+
i = torch.device("cuda:1")
854+
assert comparator(g, h)
855+
assert not comparator(g, i)
856+
857+
# Test cuda device with and without explicit index
858+
if torch.cuda.is_available():
859+
j = torch.device("cuda:0")
860+
k = torch.device("cuda", 0)
861+
assert comparator(j, k)
862+
863+
# Test meta device
864+
l = torch.device("meta")
865+
m = torch.device("meta")
866+
n = torch.device("cpu")
867+
assert comparator(l, m)
868+
assert not comparator(l, n)
869+
870+
827871
def test_jax():
828872
try:
829873
import jax.numpy as jnp

0 commit comments

Comments
 (0)