Bug fix. Added assignment of sequence ranges in PairedActivationCache to ensure proper handling of activation data. This change enhances the functionality by directly linking sequence ranges to the first activation cache, improving data integrity checks during cache operations.

Author: jkminder

dictionary_learning/cache.py

@@ -766,6 +766,7 @@ def __init__(self, store_dir_1: str, store_dir_2: str, submodule_name: str = Non
         else:
             assert len(self.activation_cache_1) == len(self.activation_cache_2), f"Lengths do not match: {len(self.activation_cache_1)} != {len(self.activation_cache_2)}"  
             self._len = len(self.activation_cache_1)
+            self._sequence_ranges = self.activation_cache_1.sequence_ranges
 
         if self.activation_cache_1.tokens is not None and self.activation_cache_2.tokens is not None:
             assert torch.all(self.activation_cache_1.tokens[:self._len] == self.activation_cache_2.tokens[:self._len]), "Tokens do not match"