(self, x: torch.Tensor)
| 451 | self.normalized_shape = original.normalized_shape |
| 452 | |
| 453 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 454 | # Reshape if needed: kernel_fn expects (x, weight, bias[, eps]) |
| 455 | orig_shape = x.shape |
| 456 | if x.dim() > 2: |
| 457 | x_2d = x.reshape(-1, x.shape[-1]) |
| 458 | else: |
| 459 | x_2d = x |
| 460 | |
| 461 | try: |
| 462 | # Try full signature: kernel_fn(x, weight, bias, eps) |
| 463 | out = self.kernel_fn(x_2d, self.weight, self.bias, self.eps) |
| 464 | except TypeError: |
| 465 | try: |
| 466 | # Try without eps: kernel_fn(x, weight, bias) |
| 467 | out = self.kernel_fn(x_2d, self.weight, self.bias) |
| 468 | except TypeError: |
| 469 | # Fallback: just x |
| 470 | out = self.kernel_fn(x_2d) |
| 471 | |
| 472 | if len(orig_shape) > 2: |
| 473 | out = out.reshape(orig_shape) |
| 474 | |
| 475 | return out |
| 476 | |
| 477 | |
| 478 | class _RMSNormWrapper(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected