(self, x: torch.Tensor)
| 416 | self.bias = original.bias |
| 417 | |
| 418 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 419 | # Reshape to 2D for kernel_fn, then reshape back |
| 420 | orig_shape = x.shape |
| 421 | if x.dim() > 2: |
| 422 | x_2d = x.reshape(-1, x.shape[-1]) |
| 423 | else: |
| 424 | x_2d = x |
| 425 | |
| 426 | # kernel_fn expects (A, B) where A @ B = C |
| 427 | # For nn.Linear: output = input @ weight.T + bias |
| 428 | # So we call kernel_fn(input, weight.T) |
| 429 | weight_t = self.weight.t().contiguous() |
| 430 | out = self.kernel_fn(x_2d, weight_t) |
| 431 | |
| 432 | if self.bias is not None: |
| 433 | out = out + self.bias |
| 434 | |
| 435 | if len(orig_shape) > 2: |
| 436 | out = out.reshape(*orig_shape[:-1], out.shape[-1]) |
| 437 | |
| 438 | return out |
| 439 | |
| 440 | |
| 441 | class _LayerNormWrapper(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected