(
self,
width,
lru_width,
d_conv=4,
num_heads=1,
dtype=None,
tp_group=None,
tp_size=1,
)
| 229 | class Recurrent(Module): |
| 230 | |
| 231 | def __init__( |
| 232 | self, |
| 233 | width, |
| 234 | lru_width, |
| 235 | d_conv=4, |
| 236 | num_heads=1, |
| 237 | dtype=None, |
| 238 | tp_group=None, |
| 239 | tp_size=1, |
| 240 | ): |
| 241 | super().__init__() |
| 242 | self.width = width |
| 243 | self.lru_width = lru_width |
| 244 | self.d_conv = d_conv |
| 245 | self.dtype = dtype |
| 246 | |
| 247 | self.linear_x = ColumnLinear(self.width, |
| 248 | self.lru_width, |
| 249 | dtype=dtype, |
| 250 | tp_group=tp_group, |
| 251 | tp_size=tp_size, |
| 252 | gather_output=False) |
| 253 | self.linear_y = ColumnLinear(self.width, |
| 254 | self.lru_width, |
| 255 | bias=False, |
| 256 | dtype=dtype, |
| 257 | tp_group=tp_group, |
| 258 | tp_size=tp_size, |
| 259 | gather_output=False) |
| 260 | self.y_bias = Parameter(shape=(self.lru_width // tp_size, ), |
| 261 | dtype=dtype) |
| 262 | |
| 263 | self.conv1d = MambaConv1d(self.lru_width // tp_size, |
| 264 | self.d_conv, |
| 265 | dtype=self.dtype, |
| 266 | apply_silu=False) |
| 267 | |
| 268 | self.rg_lru = RgLru(self.lru_width, |
| 269 | num_heads=num_heads, |
| 270 | dtype=dtype, |
| 271 | tp_group=tp_group, |
| 272 | tp_size=tp_size) |
| 273 | |
| 274 | self.linear_out = RowLinear(self.lru_width, |
| 275 | self.width, |
| 276 | dtype=dtype, |
| 277 | tp_group=tp_group, |
| 278 | tp_size=tp_size) |
| 279 | |
| 280 | def forward(self, |
| 281 | hidden_states: Tensor, |
nothing calls this directly
no test coverage detected