[test] update test_past_key_values_format (#37614)

allow custom shapes
2025-04-22 11:07:34 +01:00
parent 1cd110c6cb
commit 362fa37da2
19 changed files with 134 additions and 166 deletions
--- a/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py
+++ b/tests/models/recurrent_gemma/test_modeling_recurrent_gemma.py
@@ -223,10 +223,6 @@ class RecurrentGemmaModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.Te
            config_and_inputs[0].position_embedding_type = type
            self.model_tester.create_and_check_model(*config_and_inputs)

-    @unittest.skip(reason="RecurrentGemma does not return pkv")
-    def test_past_key_values_format(self):
-        pass
-
    @unittest.skip(reason="RecurrentGemma only supports sdpa")
    def test_eager_matches_sdpa_generate(self):
        pass