| #pragma once |
|
|
| #include <torch/torch.h> |
|
|
| std::vector<at::Tensor> |
| selective_scan_fwd(const at::Tensor &u, const at::Tensor &delta, |
| const at::Tensor &A, const at::Tensor &B, const at::Tensor &C, |
| const c10::optional<at::Tensor> &D_, |
| const c10::optional<at::Tensor> &z_, |
| const c10::optional<at::Tensor> &delta_bias_, |
| bool delta_softplus); |
|
|
| std::vector<at::Tensor> |
| selective_scan_bwd(const at::Tensor &u, const at::Tensor &delta, |
| const at::Tensor &A, const at::Tensor &B, const at::Tensor &C, |
| const c10::optional<at::Tensor> &D_, |
| const c10::optional<at::Tensor> &z_, |
| const c10::optional<at::Tensor> &delta_bias_, |
| const at::Tensor &dout, |
| const c10::optional<at::Tensor> &x_, |
| const c10::optional<at::Tensor> &out_, |
| c10::optional<at::Tensor> dz_, |
| bool delta_softplus, |
| bool recompute_out_z); |
|
|
