Skip to main content

mul_complex_partial

Function mul_complex_partial 

Source
fn mul_complex_partial(
    x_bmhr: Tensor<4>,
    re_bha: Tensor<3>,
    im_bha: Tensor<3>,
    tail_bh: Tensor<2>,
    rope_dim: usize,
    rotate_pairwise: bool,
) -> Tensor<4>
Expand description

Multiply the rotation-active entries of x by the complex scalar (re, im) per pair — same pairing conventions as apply_rope_partial — and the pass-through entries by the real scalar tail.

re/im are [batch, nheads, num_rope_angles] (broadcast over the mimo_rank axis); tail is [batch, nheads].