Update rope.py

rope.py

@@ -18,43 +18,34 @@ def precompute_freqs_cis(
 
 
 def apply_rotary_emb(x, freqs_cis, position_ids, num_heads, rot_dim=None, interleave=False):
-    """
-    x:            (B, num_heads, q_len, head_dim)
-    freqs_cis:    (max_seq, rot_half, 2)   # [..., cos/sin]
-    position_ids: (q_len,) or (B, q_len) or scalar
-    num_heads:    int (unused here; kept for API compatibility)
-    rot_dim:      optional; if None we use min(D, 2*rot_half)
-    """
     B, H, T, D = x.shape
-    rot_half_from_freqs = freqs_cis.size(-2)                 # available rotary half
+    rot_half_from_freqs = freqs_cis.size(-2)
     rd = rot_dim or (rot_half_from_freqs * 2)
-    rd = min(rd, D)                                          # don't exceed head_dim
+    rd = min(rd, D)
 
-    x_rot, x_pass = x[..., :rd], x[..., rd:]                 # (B,H,T,rd), (B,H,T,D-rd)
+    x_rot, x_pass = x[..., :rd], x[..., rd:]
 
-    # Gather cos/sin for each position; result (B,T,rot_half_from_freqs)
+    # gather cos/sin for these positions
     if torch.is_tensor(position_ids):
-        if position_ids.dim() == 2 and position_ids.size(0) == B:         # (B,T)
-            freq = freqs_cis[position_ids]                                # (B,T,rot_half,2)
-        elif position_ids.dim() == 1 and position_ids.size(0) == T:       # (T,)
+        if position_ids.dim() == 2 and position_ids.size(0) == B:
+            freq = freqs_cis[position_ids]                                  # (B,T,rot_half,2)
+        elif position_ids.dim() == 1 and position_ids.size(0) == T:
             freq = freqs_cis[position_ids].unsqueeze(0).expand(B, -1, -1, -1)
-        else:                                                              # scalar
+        else:                                                                # scalar
             pid = position_ids.view(()).long()
             freq = freqs_cis[pid].unsqueeze(0).expand(B, T, -1, -1)
     else:
         pid = torch.tensor(position_ids, device=x.device, dtype=torch.long)
         freq = freqs_cis[pid].unsqueeze(0).expand(B, T, -1, -1)
 
-    # Trim freqs to rd//2 if needed
     rot_half = rd // 2
-    cos = freq[..., 0][..., :rot_half].unsqueeze(1)          # (B,1,T,rot_half)
-    sin = freq[..., 1][..., :rot_half].unsqueeze(1)          # (B,1,T,rot_half)
+    cos = freq[..., 0][..., :rot_half].unsqueeze(1)                          # (B,1,T,rot_half)
+    sin = freq[..., 1][..., :rot_half].unsqueeze(1)                          # (B,1,T,rot_half)
 
-    # Split real/imag and apply rotation
     x_rot = x_rot.view(B, H, T, rot_half, 2)
-    xr, xi = x_rot[..., 0], x_rot[..., 1]                    # (B,H,T,rot_half)
+    xr, xi = x_rot[..., 0], x_rot[..., 1]
     yr = xr * cos - xi * sin
     yi = xr * sin + xi * cos
-    y = torch.stack((yr, yi), dim=-1).flatten(-2)            # (B,H,T,rd)
+    y = torch.stack((yr, yi), dim=-1).flatten(-2)
 
-    return torch.cat([y.to(x.dtype), x_pass], dim=-1)
+    return torch.cat([y.to(x.dtype), x_pass], dim=-1)