Update tico/utils/validate_args_kwargs.py

Author: Sanggyu Lee

tico/serialize/operators/op_attention.py

@@ -27,7 +27,7 @@
 from tico.serialize.operators.hashable_opcode import OpCode
 from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
 from tico.serialize.operators.utils import create_builtin_operator, get_op_index
-
+from tico.utils.validate_args_kwargs import CircleAttentionArgs
 
 
 def llama_attention_forward_adapter(
@@ -78,20 +78,7 @@ def define_node(
         self,
         node: torch.fx.Node,
     ) -> circle.Operator.OperatorT:
-        (
-            hidden_states,
-            wq,
-            wk,
-            wv,
-            wo,
-            position_cos,
-            position_sin,
-            attention_mask,
-            past_key,
-            past_value,
-            cache_position,
-        ) = node.args
-
+        args = CircleAttentionArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         op_index = get_op_index(
             circle.BuiltinOperator.BuiltinOperator.ATTENTION, self._op_codes
         )

tico/utils/validate_args_kwargs.py

@@ -171,6 +171,34 @@ class CatArgs:
     dim: int = 0
 
 
+@enforce_type
+@dataclass
+class CircleAttentionArgs:
+    """
+    For circle.BuiltinOperator.BuiltinOperator.RMS_NORM
+    """
+
+
+@enforce_type
+@dataclass
+class CircleAttentionArgs:
+    """
+    For circle.BuiltinOperator.BuiltinOperator.ATTENTION
+    """
+
+    hidden_states: torch.fx.Node
+    wq: torch.fx.Node
+    wk: torch.fx.Node
+    wv: torch.fx.Node
+    wo: torch.fx.Node
+    position_cos: torch.fx.Node
+    position_sin: torch.fx.Node
+    attention_mask: torch.fx.Node
+    past_key: torch.fx.Node
+    past_value: torch.fx.Node
+    cache_position: torch.fx.Node
+
+
 @enforce_type
 @dataclass
 class CircleRMSNormArgs: