Address code review feedback: add validation for perfect square frame_tokens

Co-authored-by: anxiangsir <31175974+anxiangsir@users.noreply.github.com>

Author: Copilot

eval_encoder/attentive_probe.py

@@ -204,6 +204,9 @@ def video_to_images(videos: torch.Tensor) -> torch.Tensor:
                     # ===> Compute patch_positions for RoPE with temporal scaling to [0, 64) <===
                     # Calculate spatial grid dimensions (assume square patches)
                     patches_per_side = int(math.sqrt(frame_tokens))  # e.g., 14 for 196 tokens
+                    assert patches_per_side * patches_per_side == frame_tokens, (
+                        f"frame_tokens must be a perfect square, got {frame_tokens}"
+                    )
                     seq_len = frame_indices.shape[1]  # Number of frames sampled
 
                     # Temporal positions: use interpolated_indices (already in [0, target_frames-1])

eval_encoder/attentive_probe_codec.py

@@ -244,6 +244,9 @@ def video_to_images(videos: torch.Tensor) -> torch.Tensor:
                     # ===> Compute patch_positions for RoPE with temporal scaling to [0, 64) <===
                     # Calculate spatial grid dimensions (assume square patches)
                     patches_per_side = int(math.sqrt(frame_tokens))  # e.g., 14 for 196 tokens
+                    assert patches_per_side * patches_per_side == frame_tokens, (
+                        f"frame_tokens must be a perfect square, got {frame_tokens}"
+                    )
 
                     # Temporal positions: use interpolated_indices (already in [0, target_frames-1])
                     # Shape: [bs, seq_len] -> expand to [bs, seq_len * frame_tokens]

tests/test_onevision_encoder.py

@@ -159,9 +159,9 @@ def test_forward_from_positions_temporal_scaling(self):
         patches_per_frame = 16  # 4x4 spatial patches
         target_frames = 64
 
-        # Simulate interpolated indices in [0, 63] range
-        # For 8 sampled frames from a video, spread across 64 target frames
-        interpolated_t = torch.tensor([0, 9, 18, 27, 36, 45, 54, 63], device=device)  # [num_frames]
+        # Simulate interpolated indices in [0, target_frames-1] range
+        # For num_frames sampled frames from a video, spread across target_frames
+        interpolated_t = torch.linspace(0, target_frames - 1, num_frames, dtype=torch.long, device=device)
 
         # Spatial positions for each frame (4x4 grid)
         h_ids = torch.arange(4, device=device).repeat_interleave(4)  # [0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3]