flux_apply_rotary_emb

Apply FLUX rotary embeddings to a per-head tensor

Description

Rotates adjacent element pairs of the last dimension: out = x * cos + rotate_half(x) * sin with pairs interleaved (elements 1,2 form the first complex pair). Math in float32, result cast back to the input dtype. Reference: apply_rotary_emb with use_real_unbind_dim = -1.

Usage

flux_apply_rotary_emb(x, freqs)

Arguments

  • x: Tensor of shape [B, H, S, D] (per-head layout).
  • freqs: List of two tensors (cos, sin), each [S, D], from flux_pos_embed.

Value

Tensor with the same shape and dtype as x.