Dynamically import a kernel .py file and return the module.
(path: str)
| 394 | |
| 395 | |
| 396 | def load_kernel_module(path: str) -> Any: |
| 397 | """Dynamically import a kernel .py file and return the module.""" |
| 398 | path = os.path.abspath(path) |
| 399 | module_name = f"opt_kernel_{os.path.basename(path).replace('.py', '')}" |
| 400 | spec = importlib.util.spec_from_file_location(module_name, path) |
| 401 | if spec is None or spec.loader is None: |
| 402 | raise ImportError(f"Cannot load kernel from: {path}") |
| 403 | mod = importlib.util.module_from_spec(spec) |
| 404 | spec.loader.exec_module(mod) |
| 405 | return mod |
| 406 | |
| 407 | |
| 408 | class _LinearWrapper(nn.Module): |