Cleanup attr ref representation

Signed-off-by: Ganesan Ramalingam <grama@microsoft.com>

Author: gramalingam

onnxscript/_internal/converter.py

@@ -339,22 +339,20 @@ def tensor_name_generator() -> str:
     def _to_onnx_attr_ref(
         self, val: values.AttrRef, info: Optional[sourceinfo.SourceInfo]
     ) -> ir.Attr:
-        pytype = val.typeinfo
-        attrtype = ta.pytype_to_attrtype(pytype)
+        attrtype = val.value.type
         attrname = None
-        if attrtype is onnx.AttributeProto.FLOAT:
+        if attrtype is ir.AttributeType.FLOAT:  # onnx.AttributeProto.FLOAT:
             attrname = "value_float"
-        elif attrtype is onnx.AttributeProto.INT:
+        elif attrtype is ir.AttributeType.INT:
             attrname = "value_int"
-        elif attrtype is onnx.AttributeProto.STRING:
+        elif attrtype is ir.AttributeType.STRING:
             attrname = "value_string"
-        elif attrtype is onnx.AttributeProto.INTS:
+        elif attrtype is ir.AttributeType.INTS:
             attrname = "value_ints"
         else:
-            msg = f"Unsupported attribute type {pytype!r}."
+            msg = f"Unsupported attribute type {attrtype!r}."
             fail(info.msg(msg) if info else msg)
-        attr_type = ir.AttributeType(ta.pytype_to_attrtype(pytype))
-        return ir.Attr(attrname, attr_type, value=None, ref_attr_name=val.value.name)
+        return ir.Attr(attrname, attrtype, value=None, ref_attr_name=val.value.name)
 
     def _to_onnx_var(
         self,
@@ -369,7 +367,7 @@ def _to_onnx_var(
             result = self.emit(
                 [result_name], values.Op(self.default_opset, "Constant"), [], [attr]
             )
-            if ta.base_type_is_bool(val.typeinfo):
+            if val.as_bool:
                 # ONNX attributes use an int-encoding for bools, but ONNX tensor types
                 # distinguish between int and bool. So we cast the int tensor to a bool tensor,
                 # to promote a (python) bool attribute to a ONNX bool tensor.
@@ -1474,7 +1472,8 @@ def _translate_function_signature_common(
                 attribute_type = ta.pytype_to_attrtype(typeinfo)
                 attr = ir.Attr(x.arg, ir.AttributeType(attribute_type), default_value, None)
                 self._current_fn.append_parameter(attr)
-                self._bind(x.arg, values.AttrRef(attr, typeinfo, self._source_of(x)))
+                as_bool = ta.base_type_is_bool(typeinfo)
+                self._bind(x.arg, values.AttrRef(attr, as_bool, self._source_of(x)))
             else:
                 onnx_parameter = make_value(x.arg, typeinfo, self._source_of(x))
                 self._current_fn.append_parameter(onnx_parameter)

onnxscript/_internal/values.py

@@ -13,7 +13,6 @@
 import typing
 from enum import IntFlag
 from typing import (  # type: ignore[attr-defined]
-    TYPE_CHECKING,
     Any,
     Callable,
     ClassVar,
@@ -22,7 +21,6 @@
     Protocol,
     Sequence,
     TypeVar,
-    _GenericAlias,
 )
 
 import onnx
@@ -36,9 +34,6 @@
 from onnxscript.ir import _schemas
 from onnxscript.onnx_types import ONNXType
 
-if TYPE_CHECKING:
-    from onnxscript._internal.type_annotation import TypeAnnotationValue
-
 _R = TypeVar("_R")
 _P = ParamSpec("_P")
 
@@ -868,23 +863,16 @@ def __init__(self, value: Any, info: sourceinfo.SourceInfo) -> None:
 
 
 class AttrRef(SymbolValue):
-    def __init__(
-        self, attr: ir.Attr, typeinfo: _GenericAlias, info: sourceinfo.SourceInfo
-    ) -> None:
+    def __init__(self, attr: ir.Attr, as_bool: bool, info: sourceinfo.SourceInfo) -> None:
         """Initializes AttrRef.
 
         Arguments:
             attr: An ir.Attr representing the attribute-parameter
-            typeinfo: type annotation of the attribute.
-                op's attributes in ONNX are usually single type or list of single type.
+            as_bool: Whether the attribute is to be interpreted as a bool type (represented as int in ONNX)
             info: for debugging use.
         """
         super().__init__(attr, info)
-        self.typeinfo = typeinfo
-        if not isinstance(typeinfo, (type, _GenericAlias)):
-            # typing._GenericAlias for List[int] and List[str], etc.
-            raise TypeError(f"Expecting a type not f{type(typeinfo)} for typeinfo.")
-        self.typeinfo = typeinfo
+        self.as_bool = as_bool
 
 
 class DynamicKind(IntFlag):
@@ -901,7 +889,7 @@ def __init__(
         ir_value: ir.Value,
         kind: DynamicKind,
         info: sourceinfo.SourceInfo,
-        typeinfo: TypeAnnotationValue | None = None,
+        typeinfo: type_annotation.TypeAnnotationValue | None = None,
     ) -> None:
         """Represents an ir.Value with some extra information.