(chunksize=64)
| 379 | |
| 380 | # api to fit original mamba_ssm |
| 381 | def build_api_selective_scan(chunksize=64): |
| 382 | mode = MODE |
| 383 | |
| 384 | def selective_scan_fn(u, delta, A, B, C, D, z=None, |
| 385 | delta_bias=None, delta_softplus=None, |
| 386 | return_last_state=False): |
| 387 | assert z is None |
| 388 | if mode in ["fnDEBUG"]: |
| 389 | outs = SelectiveScanEasy.forward(SelectiveScanEasy, u, delta, A, B, C, D, delta_bias, delta_softplus, return_last_state, chunksize) |
| 390 | return (outs[0].to(u.dtype), *outs[1:]) if return_last_state else outs[0].to(u.dtype) |
| 391 | elif mode in ["fn"]: |
| 392 | outs = SelectiveScanEasy.apply(u, delta, A, B, C, D, delta_bias, delta_softplus, return_last_state, 1) |
| 393 | return (outs[0].to(u.dtype), *outs[1:]) if return_last_state else outs[0].to(u.dtype) |
| 394 | elif mode in ["v0"]: |
| 395 | return selective_scan_easy(u, delta, A, B, C, D, delta_bias, delta_softplus, return_last_state, chunksize) |
| 396 | else: |
| 397 | raise NotImplementedError |
| 398 | return selective_scan_fn |
| 399 | |
| 400 | |
| 401 | def selective_scan_ref(u, delta, A, B, C, D=None, z=None, delta_bias=None, delta_softplus=False, |
no outgoing calls
no test coverage detected