Remove the RefAttr class (#2328)

## Rational

We defined the class `RefAttr` in the IR to represent reference
attributes in ONNX. Node attributes can be `Attr` and `RefAttr`.
However, since most of the time we are working with concrete attributes,
the union of types creates a typing situation where we always need to
assert the types before taking the values, even if we know a `RefAttr`
cannot exist (outside of a function definition).

This additionally matches the definition of AttributeProto in ONNX.

## Change

This change merged the two classes, and instead defines a `is_ref()`
method for users to check the reference attribute.

The change is BC breaking for usage like `isinstance(attr, ir.RefAttr)`.
Fortunately all such usages exist in this code base and not in PyTorch,
so we are safe to complete the change.

Author: justinchuby

onnxscript/ir/_convenience/__init__.py

@@ -32,7 +32,6 @@
     _protocols.TensorProtocol,  # This includes all in-memory tensor types
     onnx.TensorProto,
     _core.Attr,
-    _core.RefAttr,
     _protocols.GraphProtocol,
     Sequence[_protocols.GraphProtocol],
     onnx.GraphProto,
@@ -50,7 +49,7 @@ def _infer_attribute_type(attr: SupportedAttrTypes) -> _enums.AttributeType:
         return _enums.AttributeType.FLOAT
     if isinstance(attr, str):
         return _enums.AttributeType.STRING
-    if isinstance(attr, (_core.Attr, _core.RefAttr)):
+    if isinstance(attr, _core.Attr):
         return attr.type
     if isinstance(attr, Sequence) and all(isinstance(x, int) for x in attr):
         return _enums.AttributeType.INTS
@@ -97,7 +96,7 @@ def convert_attribute(
     name: str,
     attr: SupportedAttrTypes,
     attr_type: _enums.AttributeType | None = None,
-) -> _core.Attr | _core.RefAttr:
+) -> _core.Attr:
     """Convert a Python object to a _core.Attr object.
 
     This method is useful when constructing nodes with attributes. It infers the
@@ -121,7 +120,7 @@ def convert_attribute(
             raise ValueError("attr_type must be provided when attr is None")
         return _core.Attr(name, attr_type, None)
 
-    if isinstance(attr, (_core.Attr, _core.RefAttr)):
+    if isinstance(attr, _core.Attr):
         if attr.name != name:
             raise ValueError(
                 f"Attribute name '{attr.name}' does not match provided name '{name}'"
@@ -181,7 +180,7 @@ def convert_attribute(
 
 def convert_attributes(
     attrs: Mapping[str, SupportedAttrTypes],
-) -> list[_core.Attr | _core.RefAttr]:
+) -> list[_core.Attr]:
     """Convert a dictionary of attributes to a list of _core.Attr objects.
 
     It infers the attribute type based on the type of the value. The supported
@@ -247,7 +246,7 @@ def convert_attributes(
     Returns:
         A list of _core.Attr objects.
     """
-    attributes: list[_core.Attr | _core.RefAttr] = []
+    attributes: list[_core.Attr] = []
     for name, attr in attrs.items():
         if attr is not None:
             attributes.append(convert_attribute(name, attr))

onnxscript/ir/_convenience/_constructors.py

@@ -194,7 +194,7 @@ def node(
         A node with the given op_type and inputs.
     """
     if attributes is None:
-        attrs: Sequence[ir.Attr | ir.RefAttr] = ()
+        attrs: Sequence[ir.Attr] = ()
     else:
         attrs = _convenience.convert_attributes(attributes)
     return _core.Node(

onnxscript/ir/_core.py

@@ -1321,7 +1321,7 @@ def __init__(
         domain: str,
         op_type: str,
         inputs: Iterable[Value | None],
-        attributes: Iterable[Attr | RefAttr] = (),
+        attributes: Iterable[Attr] = (),
         *,
         overload: str = "",
         num_outputs: int | None = None,
@@ -1353,7 +1353,7 @@ def __init__(
             metadata_props: The metadata properties.
 
         Raises:
-            TypeError: If the attributes are not :class:`Attr` or :class:`RefAttr`.
+            TypeError: If the attributes are not :class:`Attr`.
             ValueError: If ``num_outputs``, when not ``None``, is not the same as the length of the outputs.
             ValueError: If an output value is ``None``, when outputs is specified.
             ValueError: If an output value has a producer set already, when outputs is specified.
@@ -1368,13 +1368,13 @@ def __init__(
         # Values belong to their defining nodes. The values list is immutable
         self._outputs: tuple[Value, ...] = self._create_outputs(num_outputs, outputs)
         attributes = tuple(attributes)
-        if attributes and not isinstance(attributes[0], (Attr, RefAttr)):
+        if attributes and not isinstance(attributes[0], Attr):
             raise TypeError(
-                f"Expected the attributes to be Attr or RefAttr, got {type(attributes[0])}. "
+                f"Expected the attributes to be Attr, got {type(attributes[0])}. "
                 "If you are copying the attributes from another node, make sure you call "
                 "node.attributes.values() because it is a dictionary."
             )
-        self._attributes: OrderedDict[str, Attr | RefAttr] = OrderedDict(
+        self._attributes: OrderedDict[str, Attr] = OrderedDict(
             (attr.name, attr) for attr in attributes
         )
         self._overload: str = overload
@@ -1633,7 +1633,7 @@ def outputs(self, _: Sequence[Value]) -> None:
         raise AttributeError("outputs is immutable. Please create a new node instead.")
 
     @property
-    def attributes(self) -> OrderedDict[str, Attr | RefAttr]:
+    def attributes(self) -> OrderedDict[str, Attr]:
         """The attributes of the node."""
         return self._attributes
 
@@ -3106,22 +3106,28 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}({self.domain!r}, {self.name!r}, {self.overload!r}, inputs={self.inputs!r}, attributes={self.attributes!r}), outputs={self.outputs!r})"
 
 
-class RefAttr(_protocols.ReferenceAttributeProtocol, _display.PrettyPrintable):
-    """Reference attribute."""
+class Attr(
+    _protocols.AttributeProtocol,
+    _protocols.ReferenceAttributeProtocol,
+    _display.PrettyPrintable,
+):
+    """Base class for ONNX attributes or references."""
 
-    __slots__ = ("_name", "_ref_attr_name", "_type", "doc_string")
+    __slots__ = ("_name", "_ref_attr_name", "_type", "_value", "doc_string")
 
     def __init__(
         self,
         name: str,
-        ref_attr_name: str,
         type: _enums.AttributeType,
+        value: Any,
+        ref_attr_name: str | None = None,
         *,
         doc_string: str | None = None,
-    ) -> None:
+    ):
         self._name = name
-        self._ref_attr_name = ref_attr_name
         self._type = type
+        self._value = value
+        self._ref_attr_name = ref_attr_name
         self.doc_string = doc_string
 
     @property
@@ -3132,43 +3138,21 @@ def name(self) -> str:
     def name(self, value: str) -> None:
         self._name = value
 
-    @property
-    def ref_attr_name(self) -> str:
-        return self._ref_attr_name
-
-    @ref_attr_name.setter
-    def ref_attr_name(self, value: str) -> None:
-        self._ref_attr_name = value
-
     @property
     def type(self) -> _enums.AttributeType:
         return self._type
 
-    @type.setter
-    def type(self, value: _enums.AttributeType) -> None:
-        self._type = value
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}({self._name!r}, {self._type!r}, ref_attr_name={self.ref_attr_name!r})"
-
-
-class Attr(_protocols.AttributeProtocol, _display.PrettyPrintable):
-    """Base class for ONNX attributes."""
+    @property
+    def value(self) -> Any:
+        return self._value
 
-    __slots__ = ("doc_string", "name", "type", "value")
+    @property
+    def ref_attr_name(self) -> str | None:
+        return self._ref_attr_name
 
-    def __init__(
-        self,
-        name: str,
-        type: _enums.AttributeType,
-        value: Any,
-        *,
-        doc_string: str | None = None,
-    ):
-        self.name = name
-        self.type = type
-        self.value = value
-        self.doc_string = doc_string
+    def is_ref(self) -> bool:
+        """Check if this attribute is a reference attribute."""
+        return self.ref_attr_name is not None
 
     def __eq__(self, other: object) -> bool:
         if not isinstance(other, _protocols.AttributeProtocol):
@@ -3185,11 +3169,15 @@ def __eq__(self, other: object) -> bool:
         return True
 
     def __str__(self) -> str:
+        if self.is_ref():
+            return f"@{self.ref_attr_name}"
         if self.type == _enums.AttributeType.GRAPH:
             return textwrap.indent("\n" + str(self.value), " " * 4)
         return str(self.value)
 
     def __repr__(self) -> str:
+        if self.is_ref():
+            return f"{self.__class__.__name__}({self.name!r}, {self.type!r}, ref_attr_name={self.ref_attr_name!r})"
         return f"{self.__class__.__name__}({self.name!r}, {self.type!r}, {self.value!r})"
 
     # Well typed getters
@@ -3269,6 +3257,29 @@ def as_graphs(self) -> Sequence[Graph]:
 
 
 # NOTE: The following functions are just for convenience
+
+
+def RefAttr(
+    name: str,
+    ref_attr_name: str,
+    type: _enums.AttributeType,
+    doc_string: str | None = None,
+) -> Attr:
+    """Create a reference attribute.
+
+    Args:
+        name: The name of the attribute.
+        type: The type of the attribute.
+        ref_attr_name: The name of the referenced attribute.
+        doc_string: Documentation string.
+
+    Returns:
+        A reference attribute.
+    """
+    # NOTE: The function name is capitalized to maintain API backward compatibility.
+    return Attr(name, type, None, ref_attr_name=ref_attr_name, doc_string=doc_string)
+
+
 def AttrFloat32(name: str, value: float, doc_string: str | None = None) -> Attr:
     """Create a float attribute."""
     # NOTE: The function name is capitalized to maintain API backward compatibility.

onnxscript/ir/_protocols.py

@@ -36,6 +36,7 @@
     Collection,
     Iterable,
     Iterator,
+    Literal,
     Mapping,
     MutableMapping,
     MutableSequence,
@@ -422,6 +423,8 @@ class AttributeProtocol(Protocol):
     value: Any
     doc_string: str | None
 
+    def is_ref(self) -> Literal[False]: ...
+
 
 @typing.runtime_checkable
 class ReferenceAttributeProtocol(Protocol):
@@ -441,6 +444,8 @@ class ReferenceAttributeProtocol(Protocol):
     type: _enums.AttributeType
     doc_string: str | None
 
+    def is_ref(self) -> Literal[True]: ...
+
 
 @typing.runtime_checkable
 class SparseTensorProtocol(Protocol):

onnxscript/ir/_tape.py

@@ -89,7 +89,7 @@ def op(
         output: ir.Value | None = None,
     ) -> ir.Value:
         if attributes is None:
-            attrs: Sequence[ir.Attr | ir.RefAttr] = ()
+            attrs: Sequence[ir.Attr] = ()
         else:
             attrs = _convenience.convert_attributes(attributes)
         output_kwargs: dict[str, Any]
@@ -141,7 +141,7 @@ def op_multi_out(
         else:
             output_kwargs = dict(outputs=outputs)
         if attributes is None:
-            attrs: Sequence[ir.Attr | ir.RefAttr] = ()
+            attrs: Sequence[ir.Attr] = ()
         else:
             attrs = _convenience.convert_attributes(attributes)
         node = ir.Node(

onnxscript/ir/external_data.py

@@ -70,7 +70,7 @@ def _all_tensors(
     # Look at constant attributes in nodes
     for node in _traversal.RecursiveGraphIterator(graph):
         for attr in node.attributes.values():
-            if isinstance(attr, _core.RefAttr):
+            if attr.is_ref():
                 continue
             if attr.type == _enums.AttributeType.TENSOR and attr.value is not None:
                 yield attr.value

onnxscript/ir/passes/common/inliner.py

@@ -52,7 +52,7 @@ class _CopyReplace:
     def __init__(
         self,
         inliner: InlinePass,
-        attr_map: dict[str, ir.Attr | ir.RefAttr],
+        attr_map: dict[str, ir.Attr],
         value_map: dict[ir.Value, ir.Value | None],
         metadata_props: dict[str, str],
         call_stack: CallStack,
@@ -83,8 +83,8 @@ def clone_optional_value(self, value: ir.Value | None) -> ir.Value | None:
             return None
         return self.clone_value(value)
 
-    def clone_attr(self, key: str, attr: ir.Attr | ir.RefAttr) -> ir.Attr | ir.RefAttr | None:
-        if isinstance(attr, ir.Attr):
+    def clone_attr(self, key: str, attr: ir.Attr) -> ir.Attr | None:
+        if not attr.is_ref():
             if attr.type == ir.AttributeType.GRAPH:
                 graph = self.clone_graph(attr.as_graph())
                 return ir.Attr(key, ir.AttributeType.GRAPH, graph, doc_string=attr.doc_string)
@@ -94,15 +94,15 @@ def clone_attr(self, key: str, attr: ir.Attr | ir.RefAttr) -> ir.Attr | ir.RefAt
                     key, ir.AttributeType.GRAPHS, graphs, doc_string=attr.doc_string
                 )
             return attr
-        assert isinstance(attr, ir.RefAttr)
+        assert attr.is_ref()
         ref_attr_name = attr.ref_attr_name
         if ref_attr_name in self._attr_map:
             ref_attr = self._attr_map[ref_attr_name]
-            if isinstance(ref_attr, ir.Attr):
+            if not ref_attr.is_ref():
                 return ir.Attr(
                     key, ref_attr.type, ref_attr.value, doc_string=ref_attr.doc_string
                 )
-            assert isinstance(ref_attr, ir.RefAttr)
+            assert ref_attr.ref_attr_name is not None
             return ir.RefAttr(
                 key, ref_attr.ref_attr_name, ref_attr.type, doc_string=ref_attr.doc_string
             )
@@ -237,7 +237,7 @@ def _instantiate_call(self, node: ir.Node, call_site_id: CallSiteId) -> NodeRepl
                 )
 
         # Identify substitutions for both inputs and attributes of the function:
-        attributes: dict[str, ir.Attr | ir.RefAttr] = node.attributes
+        attributes: dict[str, ir.Attr] = node.attributes
         default_attr_values = {
             attr.name: attr
             for attr in function.attributes.values()

onnxscript/ir/serde.py

@@ -234,9 +234,10 @@ def to_proto(ir_object: object) -> object:
         return serialize_tensor(ir_object)
     if isinstance(ir_object, _protocols.ValueProtocol):
         return serialize_value(ir_object)
-    if isinstance(ir_object, _protocols.AttributeProtocol):
+    if isinstance(ir_object, _protocols.AttributeProtocol) and not ir_object.is_ref():
         return serialize_attribute(ir_object)
     if isinstance(ir_object, _protocols.ReferenceAttributeProtocol):
+        assert ir_object.is_ref()
         return serialize_reference_attribute_into(onnx.AttributeProto(), ir_object)
     if isinstance(ir_object, _protocols.TypeProtocol):
         return serialize_type_into(onnx.TypeProto(), ir_object)
@@ -905,14 +906,14 @@ def deserialize_metadata_props(
 _deserialize_string_string_maps = deserialize_metadata_props
 
 
-def deserialize_attribute(proto: onnx.AttributeProto) -> _core.Attr | _core.RefAttr:
+def deserialize_attribute(proto: onnx.AttributeProto) -> _core.Attr:
     return _deserialize_attribute(proto, [])
 
 
 @_capture_errors(lambda proto, scoped_values: str(proto))
 def _deserialize_attribute(
     proto: onnx.AttributeProto, scoped_values: list[dict[str, _core.Value]]
-) -> _core.Attr | _core.RefAttr:
+) -> _core.Attr:
     name = proto.name
     doc_string = _get_field(proto, "doc_string")
     type_ = _enums.AttributeType(proto.type)
@@ -1465,20 +1466,10 @@ def serialize_node_into(node_proto: onnx.NodeProto, from_: _protocols.NodeProtoc
         node_proto.output.append(output.name)
 
     for attr in from_.attributes.values():
-        if isinstance(attr, _core.Attr):
+        if not attr.is_ref():
             serialize_attribute_into(node_proto.attribute.add(), from_=attr)
-        elif isinstance(attr, _core.RefAttr):
-            serialize_reference_attribute_into(node_proto.attribute.add(), from_=attr)
-        # Handle protocol attributes for completeness. We do not check them first because
-        # calling isinstance on a protocol can be slow.
-        # Most of the time, we will have Attr or RefAttr so the two branches below
-        # will not be taken.
-        elif isinstance(attr, _protocols.AttributeProtocol):
-            serialize_attribute_into(node_proto.attribute.add(), from_=attr)
-        elif isinstance(attr, _protocols.ReferenceAttributeProtocol):
-            serialize_reference_attribute_into(node_proto.attribute.add(), from_=attr)
         else:
-            raise TypeError(f"Unsupported attribute type: {type(attr)}")
+            serialize_reference_attribute_into(node_proto.attribute.add(), from_=attr)
 
 
 def serialize_tensor(tensor: _protocols.TensorProtocol) -> onnx.TensorProto: