Skip to main content

k4_ssd_state_passing

Function k4_ssd_state_passing 

Source
pub fn k4_ssd_state_passing<B: Backend>(
    intra_chunk_state_bnhpr: Tensor<B, 5>,
    da_chunk_end_bhn: Tensor<B, 3>,
    initial_state_bhpr: Tensor<B, 4>,
) -> (Tensor<B, 5>, Tensor<B, 4>)
Expand description

Based on the Kernel 4 Triton reference _state_passing_fwd_kernel (ssd_state_passing.py).

Returns:

  • chunk_input_state_bnhpr used in K5.
  • final_state_bhpr [final output].