Skip to main content

k2_ssd_bmm

burn_mamba::mamba3::double_ssd::ssd::serial

Function k2_ssd_bmm 

Source 
pub fn k2_ssd_bmm(c_bnlmhr: Tensor<6>, b_bnlmhr: Tensor<6>) -> Tensor<5>
Expand description

Compute the intra-chunk CB matrix on fused (mimo_rank-into-chunk_len) tensors.

§Arguments

  • c_bnlmhr: [batch, nchunks, chunk_len, mimo_rank, nheads, state_rank]
  • b_bnlmhr: [batch, nchunks, chunk_len, mimo_rank, nheads, state_rank]

§Returns

  • cb_bnhLMLM: [batch, nchunks, nheads, chunk_len*mimo_rank, chunk_len*mimo_rank]