diff --git a/deep_select/interface.py b/deep_select/interface.py index 06081f2..7e3d146 100644 --- a/deep_select/interface.py +++ b/deep_select/interface.py @@ -55,6 +55,15 @@ def topk( The output tensors may not be contiguous, when topk * sizeof(input.dtype or indices_dtype) is not a multiple of 32 Bytes """ + if begin is not None: + raise ValueError("`begin` is not supported now") + if hint is not None: + raise ValueError("`hint` is not supported now") + if output_idx is not None and output_idx.dtype != indices_type: + raise TypeError( + f"`output_idx` dtype ({output_idx.dtype}) must match `indices_type` ({indices_type})" + ) + N = input.shape[0] def get_empty_and_aligned_tensor(dim0: int, dim1: int, device: torch.device, dtype: torch.dtype): @@ -70,11 +79,7 @@ def get_empty_and_aligned_tensor(dim0: int, dim1: int, device: torch.device, dty output_val = get_empty_and_aligned_tensor(N, topk, device=input.device, dtype=input.dtype) if return_value else None if output_idx is None: output_idx = get_empty_and_aligned_tensor(N, topk, device=input.device, dtype=indices_type) - else: - assert output_idx.dtype == indices_type - assert begin is None, "`begin` is not supported now" - assert hint is None, "`hint` is not supported now" backend_args = ( input, topk, diff --git a/tests/test_interface_validation.py b/tests/test_interface_validation.py new file mode 100644 index 0000000..c4c9946 --- /dev/null +++ b/tests/test_interface_validation.py @@ -0,0 +1,93 @@ +"""CPU-only tests for Python interface validation. + +Run with ``python -O`` to ensure public validation does not depend on removable +assertions. The CUDA extension is replaced with a tiny fake, but real PyTorch +tensors exercise the public function's argument handling. +""" + +from __future__ import annotations + +import importlib +import sys +import types +import unittest +from pathlib import Path + +import torch + + +class InterfaceValidationTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.backend_calls = 0 + module_names = ( + "deep_select", + "deep_select.deep_select_cuda", + "deep_select.interface", + ) + cls.saved_modules = {name: sys.modules.get(name) for name in module_names} + backend = types.ModuleType("deep_select.deep_select_cuda") + backend.get_alignment_requirement = lambda: (16, 32) + + def topk(*_args: object) -> None: + cls.backend_calls += 1 + + backend.topk = topk + package = types.ModuleType("deep_select") + package.__path__ = [str(Path(__file__).resolve().parents[1] / "deep_select")] + sys.modules["deep_select"] = package + sys.modules["deep_select.deep_select_cuda"] = backend + cls.interface = importlib.import_module("deep_select.interface") + + @classmethod + def tearDownClass(cls) -> None: + for name, module in cls.saved_modules.items(): + if module is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = module + + def setUp(self) -> None: + self.__class__.backend_calls = 0 + self.input = torch.empty((1, 32), dtype=torch.float32) + + def assert_rejected_before_backend( + self, error_type: type[Exception], **kwargs: object + ) -> None: + with self.assertRaises(error_type): + self.interface.topk(self.input, 4, return_value=False, **kwargs) + self.assertEqual(self.backend_calls, 0) + + def test_begin_is_rejected_explicitly(self) -> None: + self.assert_rejected_before_backend( + ValueError, begin=torch.empty(1, dtype=torch.int32) + ) + + def test_hint_is_rejected_explicitly(self) -> None: + self.assert_rejected_before_backend( + ValueError, hint=torch.empty(1, dtype=torch.int32) + ) + + def test_preallocated_output_dtype_must_match(self) -> None: + self.assert_rejected_before_backend( + TypeError, + output_idx=torch.empty((1, 4), dtype=torch.int32), + indices_type=torch.int64, + ) + + def test_valid_preallocated_output_reaches_backend(self) -> None: + output_idx = torch.empty((1, 4), dtype=torch.int64) + output_value, returned_idx = self.interface.topk( + self.input, + 4, + return_value=False, + output_idx=output_idx, + indices_type=torch.int64, + ) + self.assertIsNone(output_value) + self.assertIs(returned_idx, output_idx) + self.assertEqual(self.backend_calls, 1) + + +if __name__ == "__main__": + unittest.main()