cleanup custom operator functionality

KLab-AI3 · Oct 20, 2024 · 535bb52 · 535bb52
1 parent 924e424
commit 535bb52
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/example/custom_operator.py b/example/custom_operator.py
@@ -28,7 +28,7 @@ def conv2d_selector(orig: torch.nn.Conv2d) -> str:
     torch_out = vgg16(input_data)
 
     ai3.swap_conv2d(
-        vgg16, conv2d_selector, None, SpecialConv)
+        vgg16, conv2d_selector, None, swap_with=SpecialConv)
     sb_out = vgg16(input_data)
     assert torch.allclose(
         torch_out, sb_out, atol=1e-4)
diff --git a/src/ai3/__init__.py b/src/ai3/__init__.py
@@ -8,8 +8,7 @@
 existing *DNN*.
 """
 
-import inspect
-from typing import Mapping, Optional, Sequence, Type, Union, TypeVar
+from typing import Mapping, Optional, Sequence, Type, Union
 from . import _core, utils, layers, _version
 from .tensor import Tensor
 
@@ -110,6 +109,7 @@ def swap_operation(
         module,
         algos: Optional[AlgorithmicSelector] = None,
         sample_input_shape: Optional[Sequence[int]] = None,
+        *,
         swap_with = None):
     """
     Swaps operations in-place out of the existing *DNN* for an implementation of
@@ -220,12 +220,13 @@ def swap_backend(module,
 def swap_conv2d(module,
                 algos: Optional['AlgorithmicSelector'] = None,
                 sample_input_shape: Optional[Sequence[int]] = None,
+                *,
                 swap_with = None):
     """
     Calls
         >>> swap_operation('conv2d', module, algos, sample_input_shape) # doctest: +SKIP
     """
-    swap_operation('conv2d', module, algos, sample_input_shape, swap_with)
+    swap_operation('conv2d', module, algos, sample_input_shape, swap_with=swap_with)
 
 
 def using_mps_and_metal() -> bool: