refactor: update tests to include device and complex parameterization for scan and lpc functions

Author: yoyolicoris

tests/test_extension.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn.functional as F
 import pytest
-from torchlpc.core import lpc_np
+from torchlpc.core import lpc_np, lpc_cuda
 
 
 from .test_grad import create_test_inputs
@@ -15,24 +15,53 @@
     "cmplx",
     [True, False],
 )
-def test_scan_cpu_equiv(samples: int, cmplx: bool):
+@pytest.mark.parametrize(
+    "device",
+    [
+        "cpu",
+        pytest.param(
+            "cuda",
+            marks=pytest.mark.skipif(
+                not torch.cuda.is_available(), reason="CUDA not available"
+            ),
+        ),
+    ],
+)
+def test_scan_equiv(samples: int, cmplx: bool, device: str):
     batch_size = 4
     x = torch.randn(
-        batch_size, samples, dtype=torch.float32 if not cmplx else torch.complex64
+        batch_size,
+        samples,
+        dtype=torch.float32 if not cmplx else torch.complex64,
+        device=device,
     )
-    A = torch.rand_like(x) * 1.8 - 0.9
-    zi = torch.randn(batch_size, dtype=x.dtype)
-
-    numba_y = torch.from_numpy(
-        lpc_np(
-            x.cpu().numpy(),
-            -A.cpu().unsqueeze(2).numpy(),
-            zi.cpu().unsqueeze(1).numpy(),
+    if cmplx:
+        A = torch.rand(
+            batch_size, samples, dtype=x.dtype, device=device
+        ).sqrt() * torch.exp(
+            2j
+            * torch.rand(batch_size, samples, dtype=x.dtype, device=device)
+            * torch.pi
         )
-    )
-    ext_y = torch.ops.torchlpc.scan_cpu(x, A, zi)
+    else:
+        A = torch.rand_like(x) * 1.8 - 0.9
+    zi = torch.randn(batch_size, dtype=x.dtype, device=device)
 
-    assert torch.allclose(numba_y, ext_y)
+    if device == "cuda":
+        numba_y = lpc_cuda(x, -A.unsqueeze(2), zi.unsqueeze(1))
+    else:
+        numba_y = torch.from_numpy(
+            lpc_np(
+                x.cpu().numpy(),
+                -A.cpu().unsqueeze(2).numpy(),
+                zi.cpu().unsqueeze(1).numpy(),
+            )
+        )
+    ext_y = torch.ops.torchlpc.scan(x, A, zi)
+
+    assert torch.allclose(numba_y, ext_y, atol=5e-7), torch.max(
+        torch.abs(numba_y - ext_y)
+    ).item()
 
 
 @pytest.mark.parametrize(
@@ -43,12 +72,12 @@ def test_scan_cpu_equiv(samples: int, cmplx: bool):
     "cmplx",
     [True, False],
 )
-def test_lpc_cpu_equiv(samples: int, cmplx: bool):
+def test_lpc_equiv(samples: int, cmplx: bool):
     batch_size = 4
     x, A, zi = tuple(
         x.to("cpu") for x in create_test_inputs(batch_size, samples, cmplx)
     )
     numba_y = torch.from_numpy(lpc_np(x.numpy(), A.numpy(), zi.numpy()))
-    ext_y = torch.ops.torchlpc.lpc_cpu(x, A, zi)
+    ext_y = torch.ops.torchlpc.lpc(x, A, zi)
 
     assert torch.allclose(numba_y, ext_y)

tests/test_grad.py

@@ -123,6 +123,10 @@ def test_float64_vs_32_cuda():
     "zi_requires_grad",
     [True, False],
 )
+@pytest.mark.parametrize(
+    "cmplx",
+    [True, False],
+)
 @pytest.mark.parametrize(
     "device",
     [
@@ -139,13 +143,25 @@ def test_parallel_scan(
     x_requires_grad: bool,
     a_requires_grad: bool,
     zi_requires_grad: bool,
+    cmplx: bool,
     device: str,
 ):
     batch_size = 2
     samples = 123
-    x = torch.randn(batch_size, samples, dtype=torch.double, device=device)
-    A = torch.rand(batch_size, samples, dtype=torch.double, device=device) * 2 - 1
-    zi = torch.randn(batch_size, dtype=torch.double, device=device)
+    dtype = torch.complex128 if cmplx else torch.double
+    x = torch.randn(batch_size, samples, dtype=dtype, device=device)
+    if cmplx:
+        A = torch.rand(
+            batch_size, samples, dtype=torch.double, device=device
+        ).sqrt() * torch.exp(
+            1j
+            * torch.rand(batch_size, samples, dtype=torch.double, device=device)
+            * 2
+            * torch.pi
+        )
+    else:
+        A = torch.rand(batch_size, samples, dtype=dtype, device=device) * 2 - 1
+    zi = torch.randn(batch_size, dtype=dtype, device=device)
 
     A.requires_grad = a_requires_grad
     x.requires_grad = x_requires_grad