fix test_rope

Author: Tingbo

model/rope.py

@@ -2,7 +2,7 @@
 from typing import Tuple
 
 
-def get_meshgrid_nd(sizes, dim=2):
+def get_meshgrid_nd(sizes):
     """
     Get n-D meshgrid with given sizes.
 
@@ -46,8 +46,8 @@ def apply_rope(
     xq_ = torch.view_as_complex(
         xq.reshape(*xq.shape[:-1], -1, 2)
     )  # [B, S, H, D//2]
-    S, H = xq_.shape[1:3]
-    freqs_cis = freqs_cis.view(1, S, H, -1)  # [S, D//2] --> [1, S, H, D//(2H)]
+    S = xq_.shape[1]
+    freqs_cis = freqs_cis.view(1, S, 1, -1)  # [S, nD//2] --> [1, S, 1, nD//2]
     # view_as_real will expand [..., D/2](complex) to [..., D/2, 2](real)
     xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3).type_as(xq)
     xk_ = torch.view_as_complex(
@@ -73,7 +73,7 @@ def get_nd_rope(
         theta: Scaling factor for frequency computation.
 
     Returns:
-        emb: Positional embedding [HW, D/2]
+        emb: Positional embedding [S, D/2]
     """
     grid = get_meshgrid_nd(sizes)  # [n, T, H, W]
 
@@ -84,10 +84,10 @@ def get_nd_rope(
             dim_list[i],
             grid[i].reshape(-1),
             theta,
-        )  # [THW, D/2]
+        )  # [THW, D_i/2]
         embs.append(emb)
 
-    emb = torch.cat(embs, dim=1)  # (THW, nD/2)
+    emb = torch.cat(embs, dim=1)  # (THW, D/2)
     return emb
 
 

tests/test_rope.py

@@ -2,7 +2,7 @@
 
 import unittest
 
-from model.rope import apply_rope, get_1d_rope, get_nd_rope
+from model.rope import apply_rope, get_1d_rope, get_nd_rope, get_meshgrid_nd
 
 
 class TestRope(unittest.TestCase):
@@ -14,11 +14,12 @@ def test_get_1d_rope(self):
         self.assertEqual(emb.shape, torch.Size([l, dim // 2]))
 
     def test_get_nd_rope(self):
-        T, H, W, dim = 8, 16, 16, 4
-        dim_list = [dim, dim, dim]
+        dim_list = [2, 2, 2]
+        dim = sum(dim_list)
+        T, H, W = 8, 16, 16
         sizes = [T, H, W]
         emb = get_nd_rope(dim_list, sizes)
-        self.assertEqual(emb.shape, torch.Size([T * H * W, 3 * dim // 2]))
+        self.assertEqual(emb.shape, torch.Size([T * H * W, dim // 2]))
 
     def test_apply_rope(self):
         B, S, H, D = 1, 16, 2, 4