v0.11.3

New features

Use kernel functions to extend layers

Up until now, it was only possible to extend existing layers with kernel layers from the Hub. Starting with this release it's also possible to extend them with kernel functions from the Hub. For instance, a silu-and-mul layer

@use_kernel_forward_from_hub("SiluAndMul")
class SiluAndMul(nn.Module):
    def forward(self, input: torch.Tensor) -> torch.Tensor:
        d = input.shape[-1] // 2
        return F.silu(input[..., :d]) * input[..., d:]

can now be extended with a silu_and_mul function from the Hub:

with use_kernel_mapping({
    "SiluAndMul": {
        "cuda": FuncRepository(
            repo_id="kernels-community/activation",
            func_name="silu_and_mul",
        ),
    }
}):
    kernelize(...)

We have added the FuncRepository, LocalFuncRepository, and LockedFuncRepository classes to load functions from regular, local, and locked repositories.

Making functions extensible

The counterpart to the previous enhancement is that functions can now also be made extensible using the new use_kernel_func_from_hub decorator:

@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 implicitly replace the function by a Torch nn.Module. Since Torch modules implement __call__, it can still be called as a function:

out = silu_and_mul(x)

However, when the function stored as part of a model/layer, it will also be kernelized:

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))

Similar to layers, the function can be kernelized using both a Hub layer and a Hub function.

What's Changed

• Split up kernels.layer into several modules by @danieldk in #187
• Add discord link to the kernel requirements doc by @danieldk in #189
• Support functions as layers by @danieldk in #188

Full Changelog: v0.11.2...v0.11.3