| 1449 | return (self.self_attention_cache.key_cache[layer_idx][0, 0].any(dim=-1)).sum() |
| 1450 | |
| 1451 | def reset(self): |
| 1452 | if hasattr(self.self_attention_cache, "reset"): |
| 1453 | self.self_attention_cache.reset() |
| 1454 | if hasattr(self.cross_attention_cache, "reset"): |
| 1455 | self.cross_attention_cache.reset() |
| 1456 | elif not hasattr(self.self_attention_cache, "reset") and not hasattr(self.cross_attention_cache, "reset"): |
| 1457 | raise ValueError( |
| 1458 | "Neither self nor cross-attention cache have valid `.reset()` methods. `.reset()` should " |
| 1459 | "only be called on compatible cache classes, such as `StaticCache` or `SlidingWindowCache`. " |
| 1460 | f"Got {self.self_attention_cache.__str__()} for the self attention cache and " |
| 1461 | f"{self.cross_attention_cache.__str__()} for the cross attention cache." |
| 1462 | ) |
| 1463 | for layer_idx in self.is_updated: |
| 1464 | self.is_updated[layer_idx] = False |
| 1465 | |
| 1466 | def reorder_cache(self, beam_idx: torch.LongTensor): |
| 1467 | """Reorders the cache for beam search, given the selected beam indices.""" |