Fix Linting errors

Author: apd10

sparse_attention_hub/sparse_attention/research_attention/maskers/fixed/implementations/oracle_top_p.py

@@ -119,7 +119,9 @@ def _compute_top_p_thresholds(
         )
         # if top_p is 1.0, then threshold_positions will be equal to sorted_scores.shape[-1]
         # which is not a valid index, so we clamp it to the last valid index
-        threshold_positions = torch.clamp(threshold_positions, max=sorted_scores.shape[-1] - 1)
+        threshold_positions = torch.clamp(
+            threshold_positions, max=sorted_scores.shape[-1] - 1
+        )
         thresholds = torch.gather(sorted_scores, dim=-1, index=threshold_positions)
         return thresholds
 

sparse_attention_hub/sparse_attention/research_attention/maskers/sampling/implementations/adaptive_sampling.py

@@ -187,13 +187,13 @@ def _compute_exp_attention_scores(
 
     def _get_sampling_range(self, seq_len_keys: int) -> tuple[int, int, int]:
         """Get sampling range and validate it.
-        
+
         Args:
             seq_len_keys: Number of keys in the sequence.
-            
+
         Returns:
             Tuple of (start_idx, end_idx, sampling_range).
-            
+
         Raises:
             ValueError: If the computed sampling range is invalid.
         """
@@ -202,13 +202,13 @@ def _get_sampling_range(self, seq_len_keys: int) -> tuple[int, int, int]:
             start_idx: int = int(self.init_offset * seq_len_keys)
         else:
             start_idx = self.init_offset
-            
+
         # Compute end index
         if isinstance(self.local_offset, float):
             end_idx: int = seq_len_keys - int(self.local_offset * seq_len_keys)
         else:
             end_idx = seq_len_keys - self.local_offset
-            
+
         sampling_range = end_idx - start_idx
 
         if sampling_range <= 0:

tests/unit/sparse_attention/research_attention/maskers/fixed/implementations/test_hashattention_top_k.py

@@ -250,7 +250,12 @@ def test_compute_hashattetion_scores(self, basic_config, test_tensors):
                     keys=test_tensors["keys"],
                     queries=test_tensors["queries"],
                     attention_mask=None,
-                    previous_dense_mask=torch.zeros(test_tensors["batch_size"], test_tensors["num_heads"], test_tensors["seq_len_queries"], test_tensors["seq_len_keys"]),
+                    previous_dense_mask=torch.zeros(
+                        test_tensors["batch_size"],
+                        test_tensors["num_heads"],
+                        test_tensors["seq_len_queries"],
+                        test_tensors["seq_len_keys"],
+                    ),
                     sparse_meta_data=sparse_meta_data,
                     layer_idx=0,
                 )

tests/unit/sparse_attention/research_attention/maskers/fixed/implementations/test_oracle_top_p.py

@@ -159,7 +159,9 @@ def max_normalized(x):
         scores = masker._compute_exp_attention_scores(
             keys,
             queries,
-            previous_dense_mask=torch.zeros(batch_size, num_heads, seq_len_queries, seq_len_keys),
+            previous_dense_mask=torch.zeros(
+                batch_size, num_heads, seq_len_queries, seq_len_keys
+            ),
             attention_mask=None,
             scaling=1.0,
         )

tests/unit/sparse_attention/research_attention/maskers/sampling/test_adaptive_sampling.py

@@ -196,7 +196,9 @@ def test_compute_exp_attention_scores(self, masker, sample_tensors):
         """Test exponential attention scores computation."""
         keys, queries, _, _ = sample_tensors
 
-        exp_scores = masker._compute_exp_attention_scores(queries, keys, scaling=1.0, attention_mask=None)
+        exp_scores = masker._compute_exp_attention_scores(
+            queries, keys, scaling=1.0, attention_mask=None
+        )
 
         assert exp_scores.shape == (2, 4, 8, 16)
         assert torch.all(exp_scores >= 0)  # Exponential should be non-negative
@@ -357,7 +359,8 @@ def test_add_mask_early_exit(self, masker, sample_tensors):
         # Create a full mask
         full_mask = Mask.create_full_mask((2, 4, 8, 16), dtype=torch.float32)
 
-        result = masker.add_mask(keys,
+        result = masker.add_mask(
+            keys,
             queries,
             values,
             attention_mask,
@@ -376,7 +379,8 @@ def test_add_mask_basic(self, masker, sample_tensors):
         # Create an empty mask
         empty_mask = Mask.create_empty_mask((2, 4, 8, 16), dtype=torch.float32)
 
-        result = masker.add_mask(keys,
+        result = masker.add_mask(
+            keys,
             queries,
             values,
             attention_mask,
@@ -420,7 +424,8 @@ def test_device_consistency(self, masker, sample_tensors):
 
         empty_mask = Mask.create_empty_mask((2, 4, 8, 16), dtype=torch.float32)
 
-        result = masker.add_mask(keys,
+        result = masker.add_mask(
+            keys,
             queries,
             values,
             attention_mask,
@@ -443,7 +448,8 @@ def test_numerical_stability(self, masker, sample_tensors):
 
         empty_mask = Mask.create_empty_mask((2, 4, 8, 16), dtype=torch.float32)
 
-        result = masker.add_mask(keys,
+        result = masker.add_mask(
+            keys,
             queries,
             values,
             attention_mask,