Simplify handling of kernelized functions (#622)

* kernels: deprecate `FuncRepository` classes

* Make `use_kernel_forward_from_hub` work with functions

Change `use_kernel_forward_from_hub` to work with functions.
`use_kernel_func_from_hub` is now deprecated and its implementation is
now just a call to `use_kernel_forward_from_hub`.

* kernels: add `use_kernelized_func` decorator from transformers

This function attaches a function that is made kernelizeable to a
module, so that it can be discovered by `kernelize`.

* Sync documentation with kernel function changes

* docs: fix markup

* doc fixes

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

* Clarify and check `use_kernelized_func` must be used with kernel funcs

* Use admonitions for warnings

* Add description `MyCustomLayers`

* Remove stale comment

* Change exception to an assertion

* Stray print

* Test for deprecation warnings

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Author: danieldk

docs/source/api/layers.md

@@ -10,6 +10,10 @@
 
 [[autodoc]] kernels.use_kernel_func_from_hub
 
+### use_kernelized_func
+
+[[autodoc]] kernels.use_kernelized_func
+
 ### replace_kernel_forward_from_hub
 
 [[autodoc]] kernels.replace_kernel_forward_from_hub

docs/source/layers.md

@@ -47,32 +47,40 @@ compatible with layers from the hub.
 
 Sometimes it can be useful to make a function extensible, for example
 because the function cannot be replaced by a layer. In such cases, you
-can annotate the function with the [`~kernels.use_kernel_func_from_hub`] decorator:
+can use the [`~kernels.use_kernel_forward_from_hub`] decorator on a
+function:
 
 ```python
-@use_kernel_func_from_hub("silu_and_mul")
+@use_kernel_forward_from_hub("silu_and_mul")
 def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
     d = x.shape[-1] // 2
     return F.silu(x[..., :d]) * x[..., d:]
 ```
 
 This will replace the function by an instantiated `torch.nn.Module`
-(singleton) that calls the function itself in its forward method.
+(singleton) that calls the function itself in its forward method. It will
+still behave as a function, since `torch.nn.Module` provides an
+implementation of `__call__` that delegates to `forward`.
 
-**Note:** for kernelization to see the function, it must be a member of
-another `torch.nn.Module` that is part of the model. For example:
+For [`~kernels.kernelize`] to see the function, you must use
+[`~kernels.use_kernelized_func`] on an `torch.nn.Module` that is part
+of the to-be kernlized model to make the function visible to
+[`~kernels.kernelize`]. The function is typically attached to the module
+that uses it. For example:
 
 ```python
+@use_kernelized_func(silu_and_mul)
 class FeedForward(nn.Module):
   def __init__(self, in_features: int, out_features: int):
       self.linear = nn.Linear(in_features, out_features)
-      # Note: silu_and_mul is a Torch module.
-      self.silu_and_mul = silu_and_mul
 
   def forward(self, x: torch.Tensor) -> torch.Tensor:
-      return self.silu_and_mul(self.linear(x))
+      return silu_and_mul(self.linear(x))
 ```
 
+Functions used by [`~kernels.use_kernelized_func`] must always have a
+[`~kernels.use_kernel_forward_from_hub`] decorator.
+
 ## Kernelizing a model
 
 A model will not use Hub kernels by default, even if it contains extensible

kernels/src/kernels/__init__.py

@@ -23,6 +23,7 @@
     use_kernel_forward_from_hub,
     use_kernel_func_from_hub,
     use_kernel_mapping,
+    use_kernelized_func,
 )
 from kernels.utils import (
     LoadedKernel,
@@ -75,4 +76,5 @@
     "use_kernel_forward_from_hub",
     "use_kernel_func_from_hub",
     "use_kernel_mapping",
+    "use_kernelized_func",
 ]

kernels/src/kernels/layer/__init__.py

@@ -16,6 +16,7 @@
     LockedLayerRepository,
     replace_kernel_forward_from_hub,
     use_kernel_forward_from_hub,
+    use_kernelized_func,
 )
 from .mode import Mode
 
@@ -36,4 +37,5 @@
     "use_kernel_forward_from_hub",
     "use_kernel_func_from_hub",
     "use_kernel_mapping",
+    "use_kernelized_func",
 ]

kernels/src/kernels/layer/func.py

@@ -1,11 +1,8 @@
 import functools
-import inspect
-from inspect import Parameter, Signature
+import warnings
 from pathlib import Path
 from types import ModuleType
-from typing import TYPE_CHECKING, Callable, Protocol, Type
-
-from kernels.layer.repos import RepositoryProtocol
+from typing import TYPE_CHECKING, Protocol, Type
 
 from .._versions import select_revision_or_version
 from ..utils import (
@@ -14,6 +11,8 @@
     get_kernel,
     get_local_kernel,
 )
+from .layer import _create_func_module, use_kernel_forward_from_hub
+from .repos import RepositoryProtocol
 
 if TYPE_CHECKING:
     from torch import nn
@@ -28,6 +27,10 @@ class FuncRepository:
     """
     Repository and name of a function for kernel mapping.
 
+    > [!WARNING]
+    > `FuncRepository` is deprecated and will be removed in kernels 0.17.
+    > Use [`~kernels.LayerRepository`] instead.
+
     Args:
         repo_id (`str`):
             The Hub repository containing the layer.
@@ -68,6 +71,12 @@ def __init__(
         version: int | None = None,
         trust_remote_code: bool | list[str] = False,
     ):
+        warnings.warn(
+            "FuncRepository is deprecated and will be removed in kernels 0.17. Use LayerRepository instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
         if revision is not None and version is not None:
             raise ValueError("Either a revision or a version must be specified, not both.")
         if revision is None and version is None:
@@ -92,7 +101,9 @@ def _resolve_revision(self) -> str:
 
     def load(self) -> Type["nn.Module"]:
         kernel = get_kernel(
-            self._repo_id, revision=self._resolve_revision(), trust_remote_code=self._trust_remote_code
+            self._repo_id,
+            revision=self._resolve_revision(),
+            trust_remote_code=self._trust_remote_code,
         )
         return _get_kernel_func(self, kernel)
 
@@ -107,7 +118,15 @@ def __eq__(self, other):
         )
 
     def __hash__(self):
-        return hash((self.func_name, self._repo_id, self._revision, self._version, self._trust_remote_code))
+        return hash(
+            (
+                self.func_name,
+                self._repo_id,
+                self._revision,
+                self._version,
+                self._trust_remote_code,
+            )
+        )
 
     def __str__(self) -> str:
         return f"`{self._repo_id}` (revision: {self._resolve_revision()}), function `{self.func_name}`"
@@ -117,6 +136,10 @@ class LocalFuncRepository:
     """
     Repository and function name from a local directory for kernel mapping.
 
+    > [!WARNING]
+    > `LocalFuncRepository` is deprecated and will be removed in kernels 0.17.
+    > Use [`~kernels.LocalLayerRepository`] instead.
+
     Args:
         repo_path (`Path`):
             The local repository containing the layer.
@@ -143,6 +166,12 @@ def __init__(
         *,
         func_name: str,
     ):
+        warnings.warn(
+            "LocalFuncRepository is deprecated and will be removed in kernels 0.17. Use LocalLayerRepository instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
         self._repo_path = repo_path
         self.func_name = func_name
 
@@ -176,6 +205,10 @@ def use_kernel_func_from_hub(func_name: str):
     kernelized, it **must** be a member of another `torch.nn.Module` that is
     part of the model (see the example).
 
+    > [!WARNING]
+    > `use_kernel_func_from_hub` is deprecated and will be removed in kernels 0.17.
+    > Use [`~kernels.use_kernel_forward_from_hub`] instead.
+
     Args:
         func_name (`str`):
             The name of the function name to use for kernel lookup in registered mappings.
@@ -210,13 +243,13 @@ def forward(self, x):
         # model = kernelize(model, mode=Mode.TRAINING | Mode.TORCH_COMPILE, device="cuda")
         ```
     """
+    warnings.warn(
+        "use_kernel_func_from_hub is deprecated and will be removed in kernels 0.17. Use [`use_kernel_forward_from_hub`] instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
 
-    def decorator(func):
-        Func = _create_func_module(func)
-        Func.kernel_layer_name = func_name
-        return Func()
-
-    return decorator
+    return use_kernel_forward_from_hub(func_name)
 
 
 class LockedFuncRepository:
@@ -225,6 +258,18 @@ class LockedFuncRepository:
 
     In contrast to `FuncRepository`, this class uses repositories that
     are locked inside a project.
+
+    > [!WARNING]
+    > `LockedFuncRepository` is deprecated and will be removed in kernels 0.17.
+    > Use [`~kernels.LockedLayerRepository`] instead.
+
+    Args:
+        repo_id (`str`): The Hub repository containing the function.
+        lockfile (`Path`, *optional*): Path to the lockfile. If not provided,
+            the lockfile will be inferred from the caller's context.
+        func_name (`str`): The name of the function within the kernel repository.
+        trust_remote_code (`bool`, *optional*, defaults to `False`):
+            Whether to allow loading kernels from untrusted organisations.
     """
 
     def __init__(
@@ -238,14 +283,13 @@ def __init__(
         """
         Construct a function repository.
 
-        Args:
-            repo_id (`str`): The Hub repository containing the function.
-            lockfile (`Path`, *optional*): Path to the lockfile. If not provided,
-                the lockfile will be inferred from the caller's context.
-            func_name (`str`): The name of the function within the kernel repository.
-            trust_remote_code (`bool`, *optional*, defaults to `False`):
-                Whether to allow loading kernels from untrusted organisations.
         """
+        warnings.warn(
+            "LockedFuncRepository is deprecated and will be removed in kernels 0.17. Use LockedLayerRepository instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
         self._repo_id = repo_id
         self._lockfile = lockfile
         self.func_name = func_name
@@ -265,7 +309,11 @@ def _resolve_revision(self) -> str:
         return locked_sha
 
     def load(self) -> Type["nn.Module"]:
-        kernel = get_kernel(repo_id=self._repo_id, revision=self._revision, trust_remote_code=self._trust_remote_code)
+        kernel = get_kernel(
+            repo_id=self._repo_id,
+            revision=self._revision,
+            trust_remote_code=self._trust_remote_code,
+        )
         return _get_kernel_func(self, kernel)
 
     def __eq__(self, other):
@@ -290,27 +338,3 @@ def _get_kernel_func(repo: FuncRepositoryProtocol, kernel: ModuleType) -> Type["
         raise ValueError(f"Function `{repo.func_name}` not found in `{repo}`")
 
     return _create_func_module(func)
-
-
-def _create_func_module(func: Callable) -> Type["nn.Module"]:
-    from torch import nn
-
-    class Func(nn.Module):
-        # Same flags as possible with normal `nn.Module` objects that are exchanged
-        can_torch_compile = getattr(func, "can_torch_compile", False)
-        has_backward = getattr(func, "has_backward", True)
-
-        def forward(self, *args, **kwargs):
-            return func(*args, **kwargs)
-
-    # Use function signature with args prepended by self to support
-    # module validation.
-    func_sig = inspect.signature(func)
-    new_args = [Parameter("self", Parameter.POSITIONAL_OR_KEYWORD)]
-    new_args.extend(func_sig.parameters.values())
-    Func.forward.__signature__ = Signature(  # type: ignore[attr-defined]
-        parameters=new_args,
-        return_annotation=func_sig.return_annotation,
-    )
-
-    return Func

kernels/src/kernels/layer/kernelize.py

@@ -258,10 +258,18 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
     for _, module in model.named_modules():
         module_class = type(module)
-        if not hasattr(module_class, "kernel_layer_name"):
-            continue
 
-        kernelize_layer(module, mode=mode, device_type=device_type, use_fallback=use_fallback)
+        # Kernel functions attached to a module through `use_kernelized_func`.
+        for kernel_func in getattr(module, "_kernel_funcs", {}).values():
+            kernelize_layer(
+                kernel_func,
+                mode=mode,
+                device_type=device_type,
+                use_fallback=use_fallback,
+            )
+
+        if hasattr(module_class, "kernel_layer_name"):
+            kernelize_layer(module, mode=mode, device_type=device_type, use_fallback=use_fallback)
 
     return model