Address PR feedback

Signed-off-by: Ganesan Ramalingam <[email protected]>

Author: gramalingam

onnxscript/rewriter/_pattern_ir.py

@@ -96,7 +96,12 @@ def __str__(self) -> str:
         return self._name if self._name is not None else "anonymous:" + str(id(self))
 
 
-AttrVar = AttrPattern
+class AttrVar(AttrPattern):
+    """Represents a pattern variable used to match against attribute values."""
+
+    def __init__(self, name: str | None, *, can_match_none: bool = False):
+        super().__init__(name, can_match_none=can_match_none)
+
 
 # TODO: Support tensors. Align with usage elsewhere.
 SupportedAttrTypes = Union[
@@ -141,7 +146,7 @@ def _to_attr_pattern(value: AttrPattern | ValuePattern | SupportedAttrTypes) ->
             raise ValueError(
                 "Pattern variables used in attributes must not have check_method set."
             )
-        return AttrPattern(value.name, can_match_none=value.can_match_none)
+        return AttrVar(value.name, can_match_none=value.can_match_none)
     if isinstance(value, (int, float, str)):
         return AttrConstantPattern(value)
     if isinstance(value, Sequence):

onnxscript/rewriter/ort_fusions/mha_scale.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import math
-import onnx_ir as ir
 
 from onnxscript.rewriter import _fusion_utils, _ir_utils, pattern
 
@@ -17,20 +16,25 @@
 
 Example pattern:
     query -> Mul(scale) -> MultiHeadAttention -> output
-    
+
 Gets rewritten to:
     query -> MultiHeadAttention(with integrated scaling) -> output
 """
 
+
 class FuseMHAScale(pattern.RewriteRuleClassBase):
     def pattern(self, op, query, scale):
         scaled_query = op.Mul(query, scale)
-        mha_output = op.MultiHeadAttention(scaled_query, _allow_other_inputs=True, 
-                                           _domain="com.microsoft", _outputs=["mha_output"])
+        mha_output = op.MultiHeadAttention(
+            scaled_query,
+            _allow_other_inputs=True,
+            _domain="com.microsoft",
+            _outputs=["mha_output"],
+        )
         return mha_output
 
     def check(self, context, scale, **_):
-        scale_value =_ir_utils.get_singleton_value(scale)
+        scale_value = _ir_utils.get_singleton_value(scale)
         if scale_value is None or not isinstance(scale_value, (int, float)):
             return pattern.MatchResult().fail("Scale must be a constant numeric value.", scale)
         self._scale = scale_value
@@ -54,8 +58,11 @@ def rewrite(self, op, query, mha_output, **_):
         inputs[0] = query
         attributes = dict(attributes)
         attributes["scale"] = self._scale
-        return op.MultiHeadAttention(*inputs, **attributes, _domain="com.microsoft", _outputs=1)
+        return op.MultiHeadAttention(
+            *inputs, **attributes, _domain="com.microsoft", _outputs=1
+        )
+
 
 _mha_scale_rules = pattern.RewriteRuleSet([FuseMHAScale.rule()])
 
-fuse_mha_scale = _fusion_utils.apply_fusion_rules(_mha_scale_rules)
+fuse_mha_scale = _fusion_utils.apply_fusion_rules(_mha_scale_rules)