update tests

Author: mxmpl

benchmark/README.md

@@ -11,6 +11,11 @@ Available implementations:
 - Triton: adapter from Whisper, CUDA only.
 - PyTorch C++ extension: functions from `torchdtw`.
 
+Run with:
+```bash
+python -m dtw_benchmark
+```
+
 Computation time for DTW on array of shape (n, n), on one A40 GPU:
 
 ```

tests/conftest.py

@@ -0,0 +1,39 @@
+"""pytest configuration."""
+
+import pytest
+import torch
+from hypothesis import settings
+from hypothesis import strategies as st
+
+settings.register_profile("default", deadline=None)
+settings.load_profile("default")
+
+DIM, BATCH = st.integers(1, 1280), st.integers(1, 3)
+LOW, HIGH_MINUS_LOW = st.floats(-100, 100), st.floats(0.1, 100)
+
+
+def make_tensor(shape: tuple[int, ...], *, dtype: torch.dtype, low: float, high: float) -> torch.Tensor:
+    """Build a tensor for testing."""
+    if low == high and dtype == torch.long:
+        return torch.ones(shape, dtype=torch.long, device="cpu")
+    return torch.testing.make_tensor(shape, dtype=dtype, device="cpu", low=low, high=high)
+
+
+def assert_equal(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    """Assert tensors equal."""
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    """Add 'requires_gpu' marker."""
+    config.addinivalue_line("markers", "requires_gpu: skip test if no GPU is available")
+
+
+def pytest_collection_modifyitems(session: pytest.Session, config: pytest.Config, items: list[pytest.Item]) -> None:  # noqa: ARG001
+    """Skip tests marked with 'requires_gpu' if CUDA not available."""
+    if torch.cuda.is_available():
+        return
+    skip_gpu = pytest.mark.skip(reason="CUDA not available")
+    for item in items:
+        if "requires_gpu" in item.keywords:
+            item.add_marker(skip_gpu)

tests/test_dtw.py

@@ -2,62 +2,40 @@
 
 import pytest
 import torch
-from hypothesis import given, settings
-from hypothesis import strategies as st
+from hypothesis import given
 
 from torchdtw import dtw, dtw_batch
 
-rtol, atol = 0, 1e-9
-skipifnogpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="No GPU available")
+from .conftest import BATCH, DIM, HIGH_MINUS_LOW, LOW, assert_equal, make_tensor
 
-DIM, BATCH = st.integers(1, 1280), st.integers(1, 3)
-LOW, HIGH_MINUS_LOW = st.floats(-100, 100), st.floats(0.1, 100)
 
-
-def make_tensor(shape: tuple[int, ...], *, dtype: torch.dtype, low: float, high: float) -> torch.Tensor:
-    """Build a tensor for testing."""
-    if low == high and dtype == torch.long:
-        return torch.ones(shape, dtype=torch.long, device="cpu")
-    return torch.testing.make_tensor(shape, dtype=dtype, device="cpu", low=low, high=high)
-
-
-@skipifnogpu
+@pytest.mark.requires_gpu
 @given(x=DIM, y=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
-@settings(deadline=None)
 def test_dtw(x: int, y: int, low: float, high_minus_low: float) -> None:
     """Compare the output of dtw between CPU and GPU implementations."""
     d = make_tensor((x, y), dtype=torch.float32, low=low, high=high_minus_low + low)
-    torch.testing.assert_close(dtw(d), dtw(d.cuda()).cpu(), rtol=rtol, atol=atol)
+    assert_equal(dtw(d), dtw(d.cuda()).cpu())
 
 
-@skipifnogpu
+@pytest.mark.requires_gpu
 @given(n=BATCH, x=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
-@settings(deadline=None)
 def test_dtw_batch_symmetric(n: int, x: int, low: float, high_minus_low: float) -> None:
     """Compare the output of dtw_batch between CPU and GPU implementations, symmetric case."""
     d = make_tensor((n, n, x, x), dtype=torch.float32, low=low, high=high_minus_low + low)
     sx = make_tensor((n,), dtype=torch.long, low=1, high=x + 1)
     i, j = torch.triu_indices(n, n)
     d[i, j] = d[j, i]
-    torch.testing.assert_close(
-        dtw_batch(d, sx, sx, symmetric=True),
-        dtw_batch(d.cuda(), sx.cuda(), sx.cuda(), symmetric=True).cpu(),
-        rtol=rtol,
-        atol=atol,
-    )
+    assert_equal(dtw_batch(d, sx, sx, symmetric=True), dtw_batch(d.cuda(), sx.cuda(), sx.cuda(), symmetric=True).cpu())
 
 
-@skipifnogpu
+@pytest.mark.requires_gpu
 @given(n=BATCH, m=BATCH, x=DIM, y=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
-@settings(deadline=None)
 def test_dtw_batch_not_symmetric(n: int, m: int, x: int, y: int, low: float, high_minus_low: float) -> None:
     """Compare the output of dtw_batch between CPU and GPU implementations, non symmetric case."""
     d = make_tensor((n, m, x, y), dtype=torch.float32, low=low, high=high_minus_low + low)
     sx = make_tensor((n,), dtype=torch.long, low=1, high=x + 1)
     sy = make_tensor((m,), dtype=torch.long, low=1, high=y + 1)
-    torch.testing.assert_close(
+    assert_equal(
         dtw_batch(d, sx, sy, symmetric=False),
         dtw_batch(d.cuda(), sx.cuda(), sy.cuda(), symmetric=False).cpu(),
-        rtol=rtol,
-        atol=atol,
     )

tests/test_opcheck.py

@@ -1,54 +1,67 @@
 """Check for compatibility with torch.compile."""
 
+import pytest
 import torch
-from hypothesis import given, settings
-from hypothesis import strategies as st
+from hypothesis import given
 from torch.library import opcheck
 
 import torchdtw  # noqa: F401 # Need to import it to register dtw operation
 
-DIM, BATCH = st.integers(1, 1280), st.integers(1, 3)
-LOW, HIGH_MINUS_LOW = st.floats(-100, 100), st.floats(0.1, 100)
-CUDA_AVAILABLE = torch.cuda.is_available()
-
-
-def make_tensor(shape: tuple[int, ...], *, dtype: torch.dtype, low: float, high: float) -> torch.Tensor:
-    """Build a tensor for testing."""
-    if low == high and dtype == torch.long:
-        return torch.ones(shape, dtype=torch.long, device="cpu")
-    return torch.testing.make_tensor(shape, dtype=dtype, device="cpu", low=low, high=high)
+from .conftest import BATCH, DIM, HIGH_MINUS_LOW, LOW, make_tensor
 
 
 @given(x=DIM, y=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
-@settings(deadline=None)
 def test_opcheck_dtw(x: int, y: int, low: float, high_minus_low: float) -> None:
     """Verify that dtw can be torch compiled."""
     sample = make_tensor((x, y), dtype=torch.float32, low=low, high=high_minus_low + low)
     opcheck(torch.ops.torchdtw.dtw.default, (sample,))
-    if CUDA_AVAILABLE:
-        opcheck(torch.ops.torchdtw.dtw.default, (sample.cuda(),))
+
+
+@pytest.mark.requires_gpu
+@given(x=DIM, y=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
+def test_opcheck_dtw_cuda(x: int, y: int, low: float, high_minus_low: float) -> None:
+    """Verify that dtw can be torch compiled on CUDA."""
+    sample = make_tensor((x, y), dtype=torch.float32, low=low, high=high_minus_low + low)
+    opcheck(torch.ops.torchdtw.dtw.default, (sample.cuda(),))
 
 
 @given(n=BATCH, x=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
-@settings(deadline=None)
 def test_opcheck_dtw_batch_symmetric(n: int, x: int, low: float, high_minus_low: float) -> None:
     """Verify that dtw_batch can be torch compiled, with symmetric input."""
     sample = make_tensor((n, n, x, x), dtype=torch.float32, low=low, high=high_minus_low + low)
     sx = make_tensor((n,), dtype=torch.long, low=1, high=x)
     i, j = torch.triu_indices(n, n)
     sample[i, j] = sample[j, i]
     opcheck(torch.ops.torchdtw.dtw_batch.default, (sample, sx, sx), {"symmetric": True})
-    if CUDA_AVAILABLE:
-        opcheck(torch.ops.torchdtw.dtw_batch.default, (sample.cuda(), sx.cuda(), sx.cuda()), {"symmetric": True})
+
+
+@pytest.mark.requires_gpu
+@given(n=BATCH, x=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
+def test_opcheck_dtw_batch_symmetric_cuda(n: int, x: int, low: float, high_minus_low: float) -> None:
+    """Verify that dtw_batch can be torch compiled on CUDA, with symmetric input."""
+    sample = make_tensor((n, n, x, x), dtype=torch.float32, low=low, high=high_minus_low + low)
+    sx = make_tensor((n,), dtype=torch.long, low=1, high=x)
+    i, j = torch.triu_indices(n, n)
+    sample[i, j] = sample[j, i]
+    opcheck(torch.ops.torchdtw.dtw_batch.default, (sample.cuda(), sx.cuda(), sx.cuda()), {"symmetric": True})
 
 
 @given(n=BATCH, m=BATCH, x=DIM, y=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
-@settings(deadline=None)
 def test_opcheck_dtw_batch_not_symmetric(n: int, m: int, x: int, y: int, low: float, high_minus_low: float) -> None:
     """Verify that dtw_batch can be torch compiled, with symmetric input."""
     sample = make_tensor((n, m, x, y), dtype=torch.float32, low=low, high=high_minus_low + low)
     sx = make_tensor((n,), dtype=torch.long, low=1, high=x)
     sy = make_tensor((m,), dtype=torch.long, low=1, high=y)
     opcheck(torch.ops.torchdtw.dtw_batch.default, (sample, sx, sy), {"symmetric": False})
-    if CUDA_AVAILABLE:
-        opcheck(torch.ops.torchdtw.dtw_batch.default, (sample.cuda(), sx.cuda(), sy.cuda()), {"symmetric": False})
+
+
+@pytest.mark.requires_gpu
+@given(n=BATCH, m=BATCH, x=DIM, y=DIM, low=LOW, high_minus_low=HIGH_MINUS_LOW)
+def test_opcheck_dtw_batch_not_symmetric_cuda(
+    n: int, m: int, x: int, y: int, low: float, high_minus_low: float
+) -> None:
+    """Verify that dtw_batch can be torch compiled on CUDA, with symmetric input."""
+    sample = make_tensor((n, m, x, y), dtype=torch.float32, low=low, high=high_minus_low + low)
+    sx = make_tensor((n,), dtype=torch.long, low=1, high=x)
+    sy = make_tensor((m,), dtype=torch.long, low=1, high=y)
+    opcheck(torch.ops.torchdtw.dtw_batch.default, (sample.cuda(), sx.cuda(), sy.cuda()), {"symmetric": False})