Class _LinearWrapper

verify.py:408–438

Wraps nn.Linear to use an optimized matmul kernel_fn.

class _LinearWrapper(nn.Module):
    """Wraps nn.Linear to use an optimized matmul kernel_fn."""

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

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Reshape to 2D for kernel_fn, then reshape back
        orig_shape = x.shape
        if x.dim() > 2:
            x_2d = x.reshape(-1, x.shape[-1])
        else:
            x_2d = x

        # kernel_fn expects (A, B) where A @ B = C
        # For nn.Linear: output = input @ weight.T + bias
        # So we call kernel_fn(input, weight.T)
        weight_t = self.weight.t().contiguous()
        out = self.kernel_fn(x_2d, weight_t)

        if self.bias is not None:
            out = out + self.bias

        if len(orig_shape) > 2:
            out = out.reshape(*orig_shape[:-1], out.shape[-1])

        return out


class _LayerNormWrapper(nn.Module):