(self, x, cache)
| 180 | return output |
| 181 | |
| 182 | def step(self, x, cache): |
| 183 | # x : (B, D) |
| 184 | # cache : (h, inputs) |
| 185 | # h : (B, ED, N) |
| 186 | # inputs: (B, ED, d_conv-1) |
| 187 | |
| 188 | # output : (B, D) |
| 189 | # cache : (h, inputs) |
| 190 | |
| 191 | output, cache = self.mixer.step(self.norm(x), cache) |
| 192 | output = output + x |
| 193 | return output, cache |
| 194 | |
| 195 | |
| 196 | class MambaBlock(nn.Module): |