| 258 | self.out_proj = BitLinear(config.d_inner, config.dim, bias=config.bias) |
| 259 | |
| 260 | def forward(self, x): |
| 261 | # x : (B, L, D) |
| 262 | |
| 263 | # y : (B, L, D) |
| 264 | |
| 265 | _, L, _ = x.shape |
| 266 | |
| 267 | xz = self.in_proj(x) # (B, L, 2*ED) |
| 268 | x, z = xz.chunk(2, dim=-1) # (B, L, ED), (B, L, ED) |
| 269 | |
| 270 | # x branch |
| 271 | x = x.transpose(1, 2) # (B, ED, L) |
| 272 | x = self.conv1d(x)[ |
| 273 | :, :, :L |
| 274 | ] # depthwise convolution over time, with a short filter |
| 275 | x = x.transpose(1, 2) # (B, L, ED) |
| 276 | |
| 277 | x = F.silu(x) |
| 278 | y = self.ssm(x) |
| 279 | |
| 280 | # z branch |
| 281 | z = F.silu(z) |
| 282 | |
| 283 | output = y * z |
| 284 | output = self.out_proj(output) # (B, L, D) |
| 285 | |
| 286 | return output |
| 287 | |
| 288 | def ssm(self, x): |
| 289 | # x : (B, L, ED) |