(
self,
dim: int, # D
num_tokens: int,
sequence_length: int,
depth: int,
dt_rank: Union[int, str] = "auto",
d_state: int = 16, # N in paper/comments
expand_factor: int = 2, # E in paper/comments
d_conv: int = 4,
dt_min: float = 0.001,
dt_max: float = 0.1,
dt_init: str = "random", # "random" or "constant"
dt_scale: float = 1.0,
dt_init_floor=1e-4,
bias: bool = False,
conv_bias: bool = True,
pscan: bool = True, # use parallel scan mode or sequential mode when training
return_embeddings: bool = True,
return_tokens: bool = True,
*args,
**kwargs
)
| 565 | """ |
| 566 | |
| 567 | def __init__( |
| 568 | self, |
| 569 | dim: int, # D |
| 570 | num_tokens: int, |
| 571 | sequence_length: int, |
| 572 | depth: int, |
| 573 | dt_rank: Union[int, str] = "auto", |
| 574 | d_state: int = 16, # N in paper/comments |
| 575 | expand_factor: int = 2, # E in paper/comments |
| 576 | d_conv: int = 4, |
| 577 | dt_min: float = 0.001, |
| 578 | dt_max: float = 0.1, |
| 579 | dt_init: str = "random", # "random" or "constant" |
| 580 | dt_scale: float = 1.0, |
| 581 | dt_init_floor=1e-4, |
| 582 | bias: bool = False, |
| 583 | conv_bias: bool = True, |
| 584 | pscan: bool = True, # use parallel scan mode or sequential mode when training |
| 585 | return_embeddings: bool = True, |
| 586 | return_tokens: bool = True, |
| 587 | *args, |
| 588 | **kwargs |
| 589 | ): |
| 590 | super().__init__(*args, **kwargs) |
| 591 | self.dim = dim |
| 592 | self.num_token = num_tokens |
| 593 | self.sequence_length = sequence_length |
| 594 | self.depth = depth |
| 595 | self.dt_rank = dt_rank |
| 596 | self.d_state = d_state |
| 597 | self.expand_factor = expand_factor |
| 598 | self.d_conv = d_conv |
| 599 | self.dt_min = dt_min |
| 600 | self.dt_max = dt_max |
| 601 | self.dt_init = dt_init |
| 602 | self.dt_scale = dt_scale |
| 603 | self.dt_init_floor = dt_init_floor |
| 604 | self.bias = bias |
| 605 | self.conv_bias = conv_bias |
| 606 | self.pscan = pscan |
| 607 | self.return_embeddings = return_embeddings |
| 608 | self.return_tokens = return_tokens |
| 609 | |
| 610 | self.d_inner = self.expand_factor * self.dim |
| 611 | |
| 612 | if self.dt_rank == "auto": |
| 613 | self.dt_rank = math.ceil(self.dim / 16) |
| 614 | |
| 615 | # Mamba |
| 616 | config = MambaConfig( |
| 617 | dim=self.dim, |
| 618 | depth=self.depth, |
| 619 | dt_rank=self.dt_rank, |
| 620 | d_state=self.d_state, |
| 621 | expand_factor=self.expand_factor, |
| 622 | d_conv=self.d_conv, |
| 623 | dt_min=self.dt_min, |
| 624 | dt_max=self.dt_max, |
nothing calls this directly
no test coverage detected