[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

scripts/benchmarks/gemma4_flex_inference.py

@@ -93,7 +93,9 @@ from torch.nn.attention.flex_attention import create_block_mask as _create_block
 # (local to this file so that file is untouched).
 
 
-def _causal_blockmask_with_window(B: int, L: int, block_size: int, window: int, device: str):
+def _causal_blockmask_with_window(
+    B: int, L: int, block_size: int, window: int, device: str
+):
     def causal_windowed(b, h, q_idx, kv_idx):
         return (q_idx >= kv_idx) & (q_idx - kv_idx < window)
 
@@ -285,9 +287,7 @@ def make_flex_gemma4_attention_forward(page_table: PageTable):
     return forward
 
 
-def patch_gemma4_attention_forwards(
-    model: torch.nn.Module, page_table: PageTable
-):
+def patch_gemma4_attention_forwards(model: torch.nn.Module, page_table: PageTable):
     """Attach a PagedKVCache to every non-shared attention layer, link
     every shared attention layer to its store layer's cache, and swap in
     the flex_attention forward above.
@@ -329,9 +329,7 @@ def patch_gemma4_attention_forwards(
 # --- model forward walker --------------------------------------------------
 
 
-def call_gemma4_model_with_flex_kwargs(
-    model, input_ids, position_ids, flex_kwargs
-):
+def call_gemma4_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs):
     """Walk the Gemma-4 text model manually so we can inject flex_* kwargs
     into each attention call. Mirrors `call_model_with_flex_kwargs` in
     `qwen3_flex_inference.py` but:
@@ -594,9 +592,9 @@ class FlexGemma4Inference:
             kv_num_blocks = block_mask.kv_num_blocks[
                 batch_idx, :, input_block_idx
             ].view(B, 1, 1)
-            kv_indices = block_mask.kv_indices[
-                batch_idx, :, input_block_idx
-            ].view(B, 1, 1, -1)
+            kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view(
+                B, 1, 1, -1
+            )
             full_num = full_idx = None
             if block_mask.full_kv_num_blocks is not None:
                 full_num = block_mask.full_kv_num_blocks[
@@ -625,9 +623,7 @@ class FlexGemma4Inference:
 
         def causal_offset_windowed(off, window):
             def m(b, h, q_idx, kv_idx):
-                return (q_idx + off[b] >= kv_idx) & (
-                    q_idx + off[b] - kv_idx < window
-                )
+                return (q_idx + off[b] >= kv_idx) & (q_idx + off[b] - kv_idx < window)
 
             return m
 
@@ -696,9 +692,7 @@ class FlexGemma4Inference:
                 allocated = self.page_table.allocate()
                 self.page_table.reserve(
                     allocated,
-                    torch.tensor(
-                        [allocated], device = self.device, dtype = torch.long
-                    ),
+                    torch.tensor([allocated], device = self.device, dtype = torch.long),
                     self.page_size,
                 )
                 reserved_batches.append(allocated)
@@ -1026,9 +1020,7 @@ def main():
                 "--verify_no_drift only applies to the bf16 double-copy path."
             )
         if args.no_merge_lora:
-            raise SystemExit(
-                "--verify_no_drift is incompatible with --no_merge_lora."
-            )
+            raise SystemExit("--verify_no_drift is incompatible with --no_merge_lora.")
         if base_model is None or peft_model is None:
             raise SystemExit(
                 "--verify_no_drift requires --lora_adapter against the bf16 path."
@@ -1145,9 +1137,7 @@ def main():
         "peak_memory_gb": peak,
         "sample_completions": sample_completions,
     }
-    os.makedirs(
-        os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True
-    )
+    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
     with open(args.stats_path, "w") as f:
         json.dump(res, f, indent = 2)
     print(json.dumps(res, indent = 2))

scripts/benchmarks/verify_gemma4_numerics.py

@@ -26,7 +26,9 @@ from gemma4_flex_inference import (  # noqa: E402
 
 def main():
     Gemma4ForCausalLM, Gemma4Config, Gemma4TextConfig = _require_gemma4()
-    from transformers.models.gemma4.modeling_gemma4 import Gemma4ForConditionalGeneration
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4ForConditionalGeneration,
+    )
     from transformers import AutoTokenizer
 
     name = "unsloth/gemma-4-E2B-it"
@@ -78,9 +80,7 @@ def main():
     with torch.inference_mode():
         # `Gemma4ForConditionalGeneration.forward` applies
         # `final_logit_softcapping` internally.
-        ref_logits = ref_raw(input_ids = ids, use_cache = False).logits[
-            0, -1, :
-        ].float()
+        ref_logits = ref_raw(input_ids = ids, use_cache = False).logits[0, -1, :].float()
         shell_logits = shell(ids, use_cache = False).logits[0, -1, :].float()
     print(
         f"raw   Gemma4ForConditionalGeneration: mean {ref_logits.mean():.4f}, "