fix TestSetCosSinCache

Signed-off-by: taoyuxiang <[email protected]>

Author: NicholasTao

tests/ut/ops/test_rotary_embedding.py

@@ -1,6 +1,7 @@
 import math
 from unittest.mock import MagicMock, patch
 
+import pytest
 import torch
 
 from tests.ut.base import TestBase
@@ -316,48 +317,49 @@ def test_scale_greater_than_1(self):
                 msg=f"Failed for scale={scale}, mscale={mscale}")
 
 
-class MockRotaryEmbedding(torch.nn.Module):
+class MockRotaryEmbedding:
 
     def __init__(self, base, rotary_dim, max_position_embeddings):
-        super().__init__()
-
         self.base = base
-
         self.rotary_dim = rotary_dim
-
         self.max_position_embeddings = max_position_embeddings
 
-    def _set_cos_sin_cache(self, seq_len, device, dtype):
-        return raw__set_cos_sin_cache(self, seq_len, device, dtype)
-
-
-class TestSetCosSinCache(TestBase):
-
-    def test_set_cos_sin_cache(self):
-        # prepare an instance with reasonable values
-        base = 10000.0
-        rotary_dim = 4
-        max_pos = 10
-        model = MockRotaryEmbedding(base, rotary_dim, max_pos)
-        # mock out register_buffer
-        model.register_buffer = MagicMock()
-        # call the private method via name mangling
-        model._set_cos_sin_cache(seq_len=8, device="cpu", dtype=torch.float32)
-        # expect three calls: inv_freq, cos, sin
-        assert model.register_buffer.call_count == 3
-        names = [call.args[0] for call in model.register_buffer.call_args_list]
-        assert set(names) == {"inv_freq", "cos", "sin"}
-        # verify inv_freq shape
-        inv_freq = model.register_buffer.call_args_list[0].args[1]
-        assert isinstance(inv_freq, torch.Tensor)
-        assert inv_freq.shape == (rotary_dim // 2, )
-        # verify cos buffer
-        cos = model.register_buffer.call_args_list[1].args[1]
-        assert isinstance(cos, torch.Tensor)
-        assert cos.shape == (max_pos, rotary_dim)
-        assert cos.dtype == torch.float32
-        # verify sin buffer
-        sin = model.register_buffer.call_args_list[2].args[1]
-        assert isinstance(sin, torch.Tensor)
-        assert sin.shape == (max_pos, rotary_dim)
-        assert sin.dtype == torch.float32
+
+@pytest.fixture
+def dummy_module():
+    return MockRotaryEmbedding(base=10000.0,
+                               rotary_dim=64,
+                               max_position_embeddings=512)
+
+
+class TestSetCosSinCache:
+
+    def test_set_cos_sin_cache_generates_real_tensors(self, dummy_module):
+        calls = []
+
+        def fake_register_buffer(name, tensor, persistent=True):
+            setattr(dummy_module, name, tensor)
+            calls.append(name)
+
+        dummy_module.register_buffer = fake_register_buffer
+        seq_len = 128
+        device = torch.device("cpu")
+        dtype = torch.float32
+
+        raw__set_cos_sin_cache(dummy_module, seq_len, device, dtype)
+
+        assert calls == ['inv_freq', 'cos', 'sin']
+
+        assert isinstance(dummy_module.inv_freq, torch.Tensor)
+        assert dummy_module.inv_freq.shape == (dummy_module.rotary_dim // 2, )
+        assert dummy_module.inv_freq.device == device
+        assert dummy_module.inv_freq.dtype == torch.float32
+
+        expected_shape = (dummy_module.max_position_embeddings,
+                          dummy_module.rotary_dim)
+        for name in ('cos', 'sin'):
+            buf = getattr(dummy_module, name)
+            assert isinstance(buf, torch.Tensor)
+            assert buf.shape == expected_shape
+            assert buf.device == device
+            assert buf.dtype == torch.float32