Make _compute_dynamic_ntk_parameters exportable (#39171)

* Make _compute_dynamic_ntk_parameters exportable * add unit test
2025-07-07 14:48:31 +02:00
parent 4243bb844d
commit f16fbfb89a
2 changed files with 12 additions and 1 deletions
--- a/tests/utils/test_modeling_rope_utils.py
+++ b/tests/utils/test_modeling_rope_utils.py
@@ -220,6 +220,9 @@ class RopeTest(unittest.TestCase):
            inv_freq, _ = rope_fn(config=config, device=torch_device, seq_len=1)
            torch.testing.assert_close(inv_freq, default_inv_freq)

+            inv_freq, _ = rope_fn(config=config, device=torch_device, seq_len=torch.tensor(1, dtype=torch.int64))
+            torch.testing.assert_close(inv_freq, default_inv_freq)
+
        # Check 2: if we provide `seq_len` larger than the model's original training sequence length, the frequencies
        # will scale up (i.e., the inverse frequencies will scale down).
        factor = 10.0