| 139 | |
| 140 | @dataclass |
| 141 | class MambaConfig: |
| 142 | dim: int # D |
| 143 | depth: int |
| 144 | dt_rank: Union[int, str] = "auto" |
| 145 | d_state: int = 16 # N in paper/comments |
| 146 | expand_factor: int = 2 # E in paper/comments |
| 147 | d_conv: int = 4 |
| 148 | |
| 149 | dt_min: float = 0.001 |
| 150 | dt_max: float = 0.1 |
| 151 | dt_init: str = "random" # "random" or "constant" |
| 152 | dt_scale: float = 1.0 |
| 153 | dt_init_floor = 1e-4 |
| 154 | |
| 155 | bias: bool = False |
| 156 | conv_bias: bool = True |
| 157 | |
| 158 | pscan: bool = True # use parallel scan mode or sequential mode when training |
| 159 | |
| 160 | def __post_init__(self): |
| 161 | self.d_inner = self.expand_factor * self.dim # E*D = ED in comments |
| 162 | |
| 163 | if self.dt_rank == "auto": |
| 164 | self.dt_rank = math.ceil(self.dim / 16) |
| 165 | |
| 166 | |
| 167 | class ResidualBlock(nn.Module): |