From 7c31c2a79f604986433e5f9231211dbfc69f4bb7 Mon Sep 17 00:00:00 2001 From: Aiman Date: Fri, 20 Mar 2026 22:54:11 +0530 Subject: [PATCH] add StaticLayer.crop() to match DynamicLayer API --- src/transformers/cache_utils.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/src/transformers/cache_utils.py b/src/transformers/cache_utils.py index 7dede60a7b27..e2d868042b2b 100644 --- a/src/transformers/cache_utils.py +++ b/src/transformers/cache_utils.py @@ -353,6 +353,24 @@ def get_max_cache_shape(self) -> int: """Return the maximum cache shape of the cache""" return self.max_cache_len + def crop(self, max_length: int) -> None: + """Crop the cache to the given length.""" + if not self.is_initialized: + return + + current_length = self.cumulative_length.item() + + if max_length < 0: + raise ValueError(f"`max_length` passed to `StaticLayer.crop()` must be >= 0, got {max_length}.") + + if max_length >= current_length: + return + + self.keys[:, :, max_length:, :].zero_() + self.values[:, :, max_length:, :].zero_() + + self.cumulative_length.fill_(max_length) + class StaticSlidingWindowLayer(StaticLayer): """