Function rotate_bc_step 

pub fn rotate_bc_step(
    rot_ba: Tensor<2>,
    dt_bh: Tensor<2>,
    prev: RotationState,
    b_bmhr: Tensor<4>,
    c_bmhr: Tensor<4>,
    kind: RotationKind,
    rope_dim: usize,
) -> (Tensor<4>, Tensor<4>, RotationState)

Single-token counterpart of rotate_bc_forward for the recurrent step.

Shapes

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