diff --git a/docs/source/layers.md b/docs/source/layers.md index bf812f27..b2071cc6 100644 --- a/docs/source/layers.md +++ b/docs/source/layers.md @@ -65,10 +65,19 @@ compatible with layers from the hub. ### Using a function as a layer -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 use the [`~kernels.use_kernel_forward_from_hub`] decorator on a -function: + +> [!WARNING] +> Decorating functions is mainly intended for compatibility cases, such as in +`transformers`, where a function cannot be converted into a module. If you +control the model code, prefer converting the function to an `nn.Module` and +decorating that class with [`~kernels.use_kernel_forward_from_hub`], as shown +in [Using a decorator](#using-a-decorator). In the parent module's `__init__`, +assign an instance of the decorated class to an attribute, such as +`self.activation = SiluAndMul()`. PyTorch then registers it as a submodule, +making it part of the model graph. + +When a function cannot be replaced by a layer, you can use the +[`~kernels.use_kernel_forward_from_hub`] decorator on the function: ```python @use_kernel_forward_from_hub("silu_and_mul") diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 80dd58ae..c22050e4 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -373,6 +373,18 @@ def use_kernel_forward_from_hub(layer_name: str, condition: Callable[["nn.Module used as if it was a function. Note that a decorated function is only visible to [`~kernels.kernelize`] if it is attached to a module using [`~kernels.use_kernelized_func`]. + + + Decorating functions is mainly intended for compatibility cases, such as in + `transformers`, where a function cannot be converted into a module. If you + control the model code, prefer converting the function to an `nn.Module` and + decorating that class with [`~kernels.use_kernel_forward_from_hub`]. In the + parent module's `__init__`, assign an instance of the decorated class to an + attribute, such as `self.activation = MyActivation()`. PyTorch then registers + it as a submodule, making it part of the model graph. + + + Args: layer_name (`str`): The name of the layer to use for kernel lookup in registered mappings. @@ -484,6 +496,18 @@ def use_kernelized_func(*args: Callable): attribute (not as a submodule). This makes the function visible to [`~kernels.kernelize`]. + + + This decorator is mainly intended for compatibility cases, such as in + `transformers`, where a function cannot be converted into a module. If you + control the model code, prefer converting the function to an `nn.Module` and + decorating that class with [`~kernels.use_kernel_forward_from_hub`]. In the + parent module's `__init__`, assign an instance of the decorated class to an + attribute, such as `self.activation = MyActivation()`. PyTorch then registers + it as a submodule, making it part of the model graph. + + + Args: *args (`Callable`): Kernel functions to attach to the module.