Skip to main content

k5_ssd_chunk_scan

Function k5_ssd_chunk_scan 

Source
pub fn k5_ssd_chunk_scan<B: Backend>(
    da_cumsum_bhnl: Tensor<B, 4>,
    dt_discretized_bhnl: Tensor<B, 4>,
    x_bnlhp: Tensor<B, 5>,
    c_bnlgr: Tensor<B, 5>,
    cb_bngll: Tensor<B, 5>,
    chunk_input_state_bnhpr: Tensor<B, 5>,
    d_h: Tensor<B, 1>,
) -> Tensor<B, 5>
Expand description

Based on the Kernel 5 Triton reference _chunk_scan_fwd_kernel (ssd_chunk_scan.py).

Returns:

  • y_bnlhp [final output]