Class _LayerNormWrapper

verify.py:441–475

Wraps nn.LayerNorm to use an optimized kernel_fn.

class _LayerNormWrapper(nn.Module):
    """Wraps nn.LayerNorm to use an optimized kernel_fn."""

    def __init__(self, original: nn.LayerNorm, kernel_fn: Callable):
        super().__init__()
        self.original = original
        self.kernel_fn = kernel_fn
        self.weight = original.weight
        self.bias = original.bias
        self.eps = original.eps
        self.normalized_shape = original.normalized_shape

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Reshape if needed: kernel_fn expects (x, weight, bias[, eps])
        orig_shape = x.shape
        if x.dim() > 2:
            x_2d = x.reshape(-1, x.shape[-1])
        else:
            x_2d = x

        try:
            # Try full signature: kernel_fn(x, weight, bias, eps)
            out = self.kernel_fn(x_2d, self.weight, self.bias, self.eps)
        except TypeError:
            try:
                # Try without eps: kernel_fn(x, weight, bias)
                out = self.kernel_fn(x_2d, self.weight, self.bias)
            except TypeError:
                # Fallback: just x
                out = self.kernel_fn(x_2d)

        if len(orig_shape) > 2:
            out = out.reshape(orig_shape)

        return out


class _RMSNormWrapper(nn.Module):