Module rope 

Functions

apply_rope

Apply rotary position embeddings to x along its last dimension.

apply_rope_partial

Apply RoPE to only the rotation-active entries of the last dimension; the remainder passes through unchanged. Falls back to apply_rope when rope_dim == state_rank (full RoPE), and is the identity when rope_dim == 0 (RoPE disabled, rope_fraction = 0) — angles is ignored.

wrap_angle

Reduce angles modulo 2π into [−π, π], leaving the autodiff graph intact.