Skip to main content

rotate_bc_step

burn_mamba::mamba3::rotation

Function rotate_bc_step 

Source 
pub fn rotate_bc_step(
    rot_ba: Option<Tensor<2>>,
    dt_bh: Tensor<2>,
    prev: RotationState,
    b_bmhr: Tensor<4>,
    c_bmhr: Tensor<4>,
    spec: RotationSpec,
) -> (Tensor<4>, Tensor<4>, RotationState)
Expand description

Single-token counterpart of rotate_bc_forward for the recurrent step.

ยงShapes

  • rot_ba : [batch, num_rotation_channels], None for Real1D.
  • dt_bh : [batch, nheads].
  • b_bmhr / c_bmhr : [batch, mimo_rank, nheads, state_rank].