(
self,
num_tokens: int,
sequence_length: int,
config: MambaConfig,
return_embeddings: bool = True,
return_tokens: bool = True,
)
| 486 | |
| 487 | class Mamba(nn.Module): |
| 488 | def __init__( |
| 489 | self, |
| 490 | num_tokens: int, |
| 491 | sequence_length: int, |
| 492 | config: MambaConfig, |
| 493 | return_embeddings: bool = True, |
| 494 | return_tokens: bool = True, |
| 495 | ): |
| 496 | super().__init__() |
| 497 | self.num_tokens = num_tokens |
| 498 | self.sequence_length = sequence_length |
| 499 | self.config = config |
| 500 | self.return_embeddings = return_embeddings |
| 501 | self.return_tokens = return_tokens |
| 502 | |
| 503 | # Embedding |
| 504 | self.token_embed = nn.Embedding(num_tokens, config.dim) |
| 505 | self.norm = nn.LayerNorm(config.dim) |
| 506 | |
| 507 | self.layers = nn.ModuleList( |
| 508 | [ResidualBlock(config) for _ in range(config.depth)] |
| 509 | ) |
| 510 | # self.norm_f = RMSNorm(config.dim) |
| 511 | |
| 512 | def forward(self, x): |
| 513 | # x : (B, L, D) |
no test coverage detected