(hidden_states, weight)
| 2 | |
| 3 | @torch.no_grad() |
| 4 | def run(hidden_states, weight): |
| 5 | batch_size, hidden_size = hidden_states.shape |
| 6 | # Check constants |
| 7 | assert hidden_size == 4096 |
| 8 | |
| 9 | EPS = 1e-5 |
| 10 | |
| 11 | x = hidden_states.to(torch.float32) |
| 12 | inv_rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + EPS) |
| 13 | y = (x * inv_rms) * weight.to(torch.float32) |
| 14 | return y.to(hidden_states.dtype) |
nothing calls this directly
no outgoing calls
no test coverage detected