fix lint

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/attention.py

@@ -7,7 +7,6 @@
 import onnxscript.ir as ir
 from onnxscript.rewriter import pattern
 
-
 Dim = Union[int, ir.SymbolicDim]
 
 
@@ -43,9 +42,7 @@ def pattern(
         attention_bias,
         num_heads,
         scale,
-        **_,
     ):
-
         projected = op.MatMul(input, qkv_weight)
         # Add bias if present
         if self._has_input_bias:
@@ -71,7 +68,7 @@ def pattern(
             _allow_other_attributes=True,
             _outputs=["value_mm_sliced"],
         )
-        
+
         # Split past into past_key and past_value
         # past_key and past_value are of shape (B, H, S, D/H)
         past_key, past_value = op.Split(past, axis=0, split=[1, 1])
@@ -87,7 +84,7 @@ def pattern(
             past_key,
             past_value,
             num_heads=num_heads,
-            scale=scale,           
+            scale=scale,
             _domain="com.microsoft",
             _outputs=3,
         )
@@ -124,7 +121,7 @@ def no_match(val: ir.Value, dims: Sequence[str]) -> bool:
             return False
         if no_match(value_mm_sliced, ["B", "S", "Dh_v"]):
             return False
-        
+
         # Ensure Dh = Dh_q + Dh_k + Dh_v
         Dh = bindings.get("Dh")
         Dh_q = bindings.get("Dh_q")
@@ -169,9 +166,7 @@ def rewrite(
         )
 
 
-attention_with_input_bias_rule = Attention.rule(
-    "attention_input_bias", has_input_bias=True
-)
+attention_with_input_bias_rule = Attention.rule("attention_input_bias", has_input_bias=True)
 attention_with_no_input_bias_rule = Attention.rule(
     "attention_no_input_bias", has_input_bias=False
 )

onnxscript/rewriter/ort_fusions/attention_basic_test.py

@@ -2,13 +2,10 @@
 # Licensed under the MIT License.
 from __future__ import annotations
 
-import math
 import unittest
 
 import numpy as np
-import onnx
 import onnxruntime as ort
-import torch
 
 import onnxscript
 from onnxscript import FLOAT, script
@@ -17,15 +14,15 @@
 
 msft_op = onnxscript.values.Opset("com.microsoft", 1)
 
-# This is a basic test that verifies that a 
+# This is a basic test that verifies that a
 # proposed expanded computation using packed matmul and ORT's MHA
 # is equivalent to ORT's Attention (for the specific configuration considered).
 
 # Simple Attention: no rotary embedding, no past key/value, no cos/sin cache
 
 
 class AttentionEquivalence(unittest.TestCase):
-    def __init__(self, *args, **kwargs):  
+    def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.batchsize = 2
         self.seqlen = 8
@@ -35,7 +32,7 @@ def __init__(self, *args, **kwargs):
         self.q_hidden_size = 160
         self.k_hidden_size = 160
         self.v_hidden_size = 180
-        #self.num_groups = self.num_heads // self.kv_num_heads
+        # self.num_groups = self.num_heads // self.kv_num_heads
 
     def random_inputs(self):
         B = self.batchsize
@@ -72,6 +69,7 @@ def expanded_model_script(self):
         Dh_q = self.q_hidden_size
         Dh_qk = self.q_hidden_size + self.k_hidden_size
         Dh_qkv = self.q_hidden_size + self.k_hidden_size + self.v_hidden_size
+
         @script()
         def attention(input, weight, bias):
             QKV_no_bias = op.MatMul(input, weight)
@@ -96,9 +94,7 @@ def to_proto(self, model_script):
         D_qkv = self.q_hidden_size + self.k_hidden_size + self.v_hidden_size
         return model_script.to_model_proto(
             input_types=(FLOAT["B", "S", D], FLOAT[D, D_qkv], FLOAT[D_qkv]),
-            output_types=(
-                FLOAT["B", "S", self.v_hidden_size],
-            ),
+            output_types=(FLOAT["B", "S", self.v_hidden_size],),
         )
 
     def test_equivalence(self):