(self, x)
| 510 | # self.norm_f = RMSNorm(config.dim) |
| 511 | |
| 512 | def forward(self, x): |
| 513 | # x : (B, L, D) |
| 514 | |
| 515 | # y : (B, L, D) |
| 516 | # Embedding |
| 517 | x = self.token_embed(x) |
| 518 | x = self.norm(x) |
| 519 | |
| 520 | for layer in self.layers: |
| 521 | x = layer(x) |
| 522 | |
| 523 | x = self.norm(x) |
| 524 | |
| 525 | # Return embeddings or Logits |
| 526 | # Return Tokens |
| 527 | if self.return_tokens: |
| 528 | x = OutputHead(self.config.dim, -1)(x) |
| 529 | return x |
| 530 | else: |
| 531 | return x |
| 532 | |
| 533 | def step(self, x, caches): |
| 534 | # x : (B, L, D) |
nothing calls this directly
no outgoing calls
no test coverage detected