x: (B, S, D) kv: (B, S, H, D)
(self, x: Tensor, kv: Optional[Tensor] = None)
| 293 | nn.init.trunc_normal_(self.ffnn.ffnn.down_proj.weight, std=xstd, a=-3 * xstd, b=3 * xstd) |
| 294 | |
| 295 | def forward(self, x: Tensor, kv: Optional[Tensor] = None) -> Tensor: |
| 296 | """ |
| 297 | x: (B, S, D) |
| 298 | kv: (B, S, H, D) |
| 299 | """ |
| 300 | h = self.attn(x, kv) |
| 301 | out = self.ffnn(h) |
| 302 | return out |
| 303 | |
| 304 | |
| 305 |
nothing calls this directly
no outgoing calls
no test coverage detected