python: fix

Author: Acly

src/visp/c-api.cpp

@@ -188,7 +188,7 @@ VISP_API char const* visp_device_description(backend_device const* d) {
 
 // models
 
-struct any_model {};
+struct any_model;
 
 VISP_API int32_t visp_model_detect_family(char const* filepath, int32_t* out_family) {
     return handle_errors([&]() {

tests/CMakeLists.txt

@@ -51,7 +51,9 @@ target_link_libraries(vision-workbench PRIVATE visioncpp ggml ${VISP_FMT_LINK})
 if(VISP_CI)
   set(PYTHON_TESTS_ARGS "--ci")
 endif()
-add_test(NAME python COMMAND uv run pytest -vs tests ${PYTHON_TESTS_ARGS})
+add_test(NAME python
+  WORKING_DIRECTORY ${CMAKE_SOURCE_DIR}
+  COMMAND uv run pytest -vs tests ${PYTHON_TESTS_ARGS})
 
 #
 # Benchmarks

tests/test_primitives.py

@@ -2,7 +2,7 @@
 import torch
 
 from . import workbench
-from .workbench import input_tensor, to_nchw, to_nhwc
+from .workbench import input_tensor, to_nchw, to_nhwc, tensors_match
 
 
 def test_linear():
@@ -13,7 +13,7 @@ def test_linear():
     result = workbench.invoke_test("linear", x, dict(weight=weight, bias=bias))
 
     expected = torch.nn.functional.linear(x, weight, bias)
-    assert torch.allclose(result, expected)
+    assert tensors_match(result, expected)
 
 
 @pytest.mark.parametrize("scenario", ["stride_1_pad_0", "stride_2_pad_1", "dilation_2_pad_2"])
@@ -48,7 +48,7 @@ def test_conv_2d_depthwise(scenario: str, memory_layout: str, batch: str, backen
     if memory_layout == "nhwc":
         result = to_nchw(result)
 
-    assert torch.allclose(result, expected)
+    assert tensors_match(result, expected)
 
 
 @pytest.mark.parametrize("scenario", ["3x3", "5x5", "stride2", "nhwc"])
@@ -76,7 +76,7 @@ def test_conv_transpose_2d(scenario: str):
     if scenario == "nhwc":
         result = to_nchw(result)
 
-    assert torch.allclose(result, expected, rtol=1e-2)
+    assert tensors_match(result, expected, rtol=1e-2)
 
 
 # def test_batch_norm_2d():
@@ -106,7 +106,7 @@ def test_layer_norm():
     result = workbench.invoke_test("layer_norm", x, dict(weight=weight, bias=bias))
 
     expected = torch.nn.functional.layer_norm(x, [dim], weight, bias, eps=1e-5)
-    assert torch.allclose(result, expected, atol=1e-6)
+    assert tensors_match(result, expected, atol=1e-6)
 
 
 @pytest.mark.parametrize("backend", ["cpu", "vulkan"])
@@ -133,7 +133,7 @@ def test_window_partition(backend: str):
 
     result = workbench.invoke_test("sam_window_partition", x, {}, backend=backend)
 
-    assert torch.allclose(result, expected)
+    assert tensors_match(result, expected)
 
 
 @pytest.mark.parametrize("shift", [(0, 2, -1, 0), (0, -2, 0, 3)])
@@ -147,7 +147,7 @@ def test_roll(shift: tuple[int, int, int, int], backend: str):
     params = dict(s0=shift[3], s1=shift[2], s2=shift[1], s3=shift[0])
     result = workbench.invoke_test("roll", x, {}, params, backend)
 
-    assert torch.allclose(result, expected)
+    assert tensors_match(result, expected)
 
 
 @pytest.mark.parametrize("mode", ["bilinear", "bicubic"])
@@ -169,4 +169,4 @@ def test_interpolate(mode: str, align_corners: bool, size: str, scale: float, ba
 
     params = dict(mode=mode, h=target[0], w=target[1], align_corners=1 if align_corners else 0)
     result = workbench.invoke_test("interpolate", x, {}, params, backend)
-    assert torch.allclose(result, expected)
+    assert tensors_match(result, expected)

tests/workbench.py

@@ -1,10 +1,11 @@
 import ctypes
-from functools import reduce
-from typing import Mapping
 import torch
 import os
-
+import platform
+from functools import reduce
+from typing import Mapping
 from pathlib import Path
+from torch import Tensor
 
 float_ptr = ctypes.POINTER(ctypes.c_float)
 
@@ -90,20 +91,40 @@ def encode_params(params: Mapping[str, str | int | float]):
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 
 root_dir = Path(__file__).parent.parent
-bin_dir = root_dir / "build" / "bin"
-
-lib = ctypes.CDLL(str(bin_dir / "vision-workbench.dll"))
-lib.visp_workbench.argtypes = [
-    ctypes.c_char_p,
-    ctypes.POINTER(RawTensor),
-    ctypes.c_int32,
-    ctypes.POINTER(RawParam),
-    ctypes.c_int32,
-    ctypes.POINTER(ctypes.POINTER(RawTensor)),
-    ctypes.POINTER(ctypes.c_int32),
-    ctypes.c_int32,
-]
-lib.visp_workbench.restype = ctypes.c_int32
+
+def _load_library():
+    system = platform.system().lower()
+    if system == "windows":
+        prefix = ""
+        suffix = ".dll"
+        libdir = "bin"
+    elif system == "darwin":
+        prefix = "lib"
+        suffix = ".dylib"
+        libdir = "lib"
+    else:  # assume Linux / Unix
+        prefix = "lib"
+        suffix = ".so"
+        libdir = "lib"
+    lib_path = root_dir / "build" / libdir / f"{prefix}vision-workbench{suffix}"
+    return ctypes.CDLL(str(lib_path))
+
+try:
+    lib = _load_library()
+
+    lib.visp_workbench.argtypes = [
+        ctypes.c_char_p,
+        ctypes.POINTER(RawTensor),
+        ctypes.c_int32,
+        ctypes.POINTER(RawParam),
+        ctypes.c_int32,
+        ctypes.POINTER(ctypes.POINTER(RawTensor)),
+        ctypes.POINTER(ctypes.c_int32),
+        ctypes.c_int32,
+    ]
+    lib.visp_workbench.restype = ctypes.c_int32
+except OSError as e:
+    print(f"Error loading vision-workbench library: {e}")
 
 
 def invoke_test(
@@ -168,11 +189,14 @@ def to_nhwc(tensor: torch.Tensor):
     return tensor.permute(0, 2, 3, 1).contiguous()
 
 
-def to_nchw(tensor: torch.Tensor):
+def to_nchw(tensor: Tensor|list[Tensor]|None):
+    assert tensor is not None
+    if isinstance(tensor, list):
+        return [t.permute(0, 3, 1, 2).contiguous() for t in tensor]
     return tensor.permute(0, 3, 1, 2).contiguous()
 
 
-def convert_to_nhwc(state: dict[str, torch.Tensor], key=""):
+def convert_to_nhwc(state: dict[str, Tensor], key=""):
     for k, v in state.items():
         is_conv = (
             v.ndim == 4
@@ -250,7 +274,25 @@ def fuse_conv_2d_batch_norm(
     return False  # no match
 
 
-def print_results(result: torch.Tensor, expected: torch.Tensor):
+def print_results(result: Tensor, expected: Tensor):
     print("\ntorch seed:", torch.initial_seed())
     print("\nresult -----", result, sep="\n")
     print("\nexpected ---", expected, sep="\n")
+
+
+def tensors_match(
+    result: Tensor | list[Tensor] | None,
+    expected: Tensor | list[Tensor],
+    rtol=1e-3,
+    atol=1e-5,
+    show=False,
+):
+    assert result is not None, "No result returned"
+    if isinstance(expected, list):
+        assert isinstance(result, list), "Result is not a list"
+        assert len(result) == len(expected), f"Expected {len(expected)} tensors, got {len(result)}"
+        return all(tensors_match(r, e, rtol, atol, show) for r, e in zip(result, expected))
+    assert isinstance(result, Tensor), "Result is not a tensor"
+    if show:
+        print_results(result, expected)
+    return torch.allclose(result, expected, rtol=rtol, atol=atol)