Repository: microsoft/onnxscript

Issue:
Support Annotated types
We would like to annotate the types with `typing.Annotated` (https://docs.python.org/3/library/typing.html#typing.Annotated) to include run time value checking for the ATen lib. Annotate should only be used for attributes.

```python
@onnxscript.script()
@atenop("aten::elu")
def Elu(
    self,
    alpha: float = 1.0,
    scale: Annotated[float, Is[lambda x: x == 1.0]] = 1.0,   # Here, just need to take the float type out
    input_scale: Annotated[float, Is[lambda x: x == 1.0]] = 1.0,
):
    # del scale
    # del input_scale
    return op.Elu(self, alpha=alpha)
```


Support for annotating Optional inputs
We would like to make this work:

```python
@torch_op("aten::clamp")
def aten_clamp(
    self: TReal, min_: Optional[TReal] = None, max_: Optional[TReal] = None
) -> TReal:
    # clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor

    if op.OptionalHasElement(min_):
        min_ = op.OptionalGetElement(min_)
        min_clamp = op.CastLike(min_, self)
    else:
        # TODO: get the correct min number for the dtype
        min_clamp = op.Constant(value_float=float("-inf"))

    if op.OptionalHasElement(max_):
        max_ = op.OptionalGetElement(max_)
        max_clamp = op.CastLike(max_, self)
    else:
        max_clamp = op.Constant(value_float=float("inf"))

    # Enforce the lower and upper bounds
    clamped = op.Max(op.Min(self, max_clamp), min_clamp)
    return clamped
```

onnxruntime complains that

E           onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph: [ONNXRuntimeError] : 10 : INVALID_GRAPH : This is an invalid model. In Node, ("", OptionalHasElement, "", -1) : () -> ("output0",) , Error Node () has input size 0 not in range [min=1, max=1].

## Reference

- #76 


Discussion:
On a different note, even the use of the nested decorator `@atenop` is not supported. And may need some investigation as to what's needed there
Do you mean if the function is wrapped it's unsupported, or simply having an additional decorator is unsupported?
Discussed elsewhere. In general, a decorator will have some non-trivial effect, so ignoring (during translation/conversion) may not be a good idea. But we could expose a way for users to indicate which decorators should be treated as "no effect" and ignored.
Question: when do you get this error message? When the script is processed, or when you try to execute it with some input, etc.?
That's during eager execution 
With which opset? The op OptionalHasElement was updated in opset 18 to make the argument optional. I believe the error is valid for earlier opsets, but if it happens with opset 18, there is an issue. Am trying to run it, but as you noted elsewhere, it looks like we need some other changes before we can use the latest ONNX and opset (I am running into the same issue with default_opset you noted elsewhere.)

Resolving patch:
diff --git a/onnxscript/converter.py b/onnxscript/converter.py
--- a/onnxscript/converter.py
+++ b/onnxscript/converter.py
@@ -290,26 +290,28 @@ def generate_unique_name(self, candidate: str = "tmp") -> str:
         self.used_vars.add(r)
         return r
 
-    def to_onnx_attr_ref(self, val: values.AttrRef):
+    def to_onnx_attr_ref(self, val: values.AttrRef, info: Optional[sourceinfo.SourceInfo]):
         pytype = val.typeinfo
+        attrtype = ta.pytype_to_attrtype(pytype)
         attrname = None
-        if pytype is float:
+        if attrtype is onnx.AttributeProto.FLOAT:
             attrname = "value_float"
-        elif pytype is int:
+        elif attrtype is onnx.AttributeProto.INT:
             attrname = "value_int"
-        elif pytype is str:
+        elif attrtype is onnx.AttributeProto.STRING:
             attrname = "value_string"
-        elif pytype is List[int]:
+        elif attrtype is onnx.AttributeProto.INTS:
             attrname = "value_ints"
         else:
-            self.fail(val, f"Unsupported attribute type {pytype!r}.")
+            msg = f"Unsupported attribute type {pytype!r}."
+            fail(info.msg(msg) if info else msg)
         return self.ir_builder.make_attr_ref(attrname, val.value, pytype)
 
-    def to_onnx_var(self, val, target=None, info=None):
+    def to_onnx_var(self, val, target=None, info: Optional[sourceinfo.SourceInfo] = None):
         if isinstance(val, values.AttrRef):
             # promote attribute to value
             result = self.generate_unique_name(target if target else "tmp")
-            attr = self.to_onnx_attr_ref(val)
+            attr = self.to_onnx_attr_ref(val, info)
             self.emit([result], values.Op(self.default_opset, "Constant"), [], [attr])
             return ConverterExpression(result, ConverterExpressionKind.CONST)
         if isinstance(val, values.Dynamic):
@@ -319,7 +321,7 @@ def to_onnx_var(self, val, target=None, info=None):
         # produce a better error message otherwise
         return self.emit_const(val, target if target else "tmp", info)
 
-    def py_var_to_onnx_var(self, py_var, info):
+    def py_var_to_onnx_var(self, py_var, info: sourceinfo.SourceInfo):
         return self.to_onnx_var(self.lookup(py_var, info), target=py_var, info=info)
 
     def emit_docstring(self, docstring):
@@ -341,7 +343,6 @@ def rename(x):
             self.bind(x, values.Dynamic(r, values.DynamicKind.Output, info))
             return r
 
-        # [ self.to_onnx_var(self.lookup(pvar)) for pvar in inputs ]
         onnx_inputs = inputs
         onnx_outputs = [rename(x) for x in outputs]
         self.emit(
diff --git a/onnxscript/type_annotation.py b/onnxscript/type_annotation.py
--- a/onnxscript/type_annotation.py
+++ b/onnxscript/type_annotation.py
@@ -13,39 +13,71 @@
 
 from onnxscript.onnx_types import TensorType
 
+# TypeAnnotationValue represents the (value of) valid type-annotations recognized
+# by ONNX Script. TODO: Flesh out a formal definition. Currently, it supports
+# * float, int, str (primitive attribute types)
+# * Sequence[float], Sequence[int], Sequence[str] (attribute types)
+# * Tensor types
+# * Sequence[Tensor] types
+# * Union of above 2
+# * TypeVars with above bounds
+# * Above types with annotation attached
+TypeAnnotationValue = typing.Any
+
+# Map from python type to corresponding ONNX AttributeProto type
 _PYTYPE_TO_ATTRTYPE_MAP = {
     float: onnx.AttributeProto.FLOAT,
     int: onnx.AttributeProto.INT,
     str: onnx.AttributeProto.STRING,
-    bool: onnx.AttributeProto.INT,
+    bool: onnx.AttributeProto.INT,  # experimental
 }
 
+# Map from python type to corresponding ONNX AttributeProto type,
+# for repeated (i.e., list of) values
 _LISTTYPE_TO_ATTRTYPE_MAP = {
     float: onnx.AttributeProto.FLOATS,
     int: onnx.AttributeProto.INTS,
     str: onnx.AttributeProto.STRINGS,
-    bool: onnx.AttributeProto.INTS,
+    bool: onnx.AttributeProto.INTS,  # experimental
 }
 
 _LIST_CONSTRUCTORS = frozenset([list, typing.List, typing.Sequence, collections.abc.Sequence])
 
 
-def pytype_to_attrtype(pytype: type) -> typing.Optional[onnx.AttributeProto.AttributeType]:
+def _remove_annotation(typeinfo: TypeAnnotationValue) -> TypeAnnotationValue:
+    """Remove Annotated wrapper if present, otherwise return typeinfo as is."""
+    if hasattr(typing, "Annotated"):
+        # Present in Python 3.9+
+        if get_origin(typeinfo) is typing.Annotated:
+            return get_args(typeinfo)[0]
+    return typeinfo
+
+
+def _is_primitive_attr_type(typeinfo: TypeAnnotationValue) -> bool:
+    return typeinfo in _PYTYPE_TO_ATTRTYPE_MAP
+
+
+def pytype_to_attrtype(
+    pytype: TypeAnnotationValue,
+) -> typing.Optional[onnx.AttributeProto.AttributeType]:
+    pytype = _remove_annotation(pytype)
     if pytype in _PYTYPE_TO_ATTRTYPE_MAP:
         return _PYTYPE_TO_ATTRTYPE_MAP[pytype]
-    if get_origin(pytype) in _LIST_CONSTRUCTORS:
-        args = get_args(pytype)
-        elt_type = args[0]
+    type_constructor = get_origin(pytype)
+    # Remove Optional wrapper if present, which is represented as an Union[..., type(None)]
+    if type_constructor is typing.Union:
+        # Filter out type(None), since typing.Optional[X] evaluates to Union[X, type(None)]
+        args = [x for x in get_args(pytype) if x is not type(None)]  # noqa: E721
+        if len(args) == 1:
+            return pytype_to_attrtype(args[0])
+    if type_constructor in _LIST_CONSTRUCTORS:
+        elt_type = get_args(pytype)[0]
         if elt_type in _LISTTYPE_TO_ATTRTYPE_MAP:
             return _LISTTYPE_TO_ATTRTYPE_MAP[elt_type]
     return None
 
 
-def is_attr_type(pytype: type):
-    return pytype_to_attrtype(pytype) is not None
-
-
-def is_value_type(typeinfo):
+def _is_tensor_type(typeinfo: TypeAnnotationValue) -> bool:
     if isinstance(typeinfo, TensorType):
         return True
     if inspect.isclass(typeinfo) and issubclass(typeinfo, TensorType):
@@ -53,8 +85,52 @@ def is_value_type(typeinfo):
     return False
 
 
-def is_valid_type(typeinfo):
-    return is_attr_type(typeinfo) or is_value_type(typeinfo)
+def is_value_type(typeinfo: TypeAnnotationValue) -> bool:
+    """Returns True if typeinfo represents a value type, False if it is an attribute type.
+    Raises ValueError if typeinfo is not a supported type annotation.
+    """
+    typeinfo = _remove_annotation(typeinfo)
+    if _is_tensor_type(typeinfo):
+        return True
+    if _is_primitive_attr_type(typeinfo):
+        return False
+    type_constructor = get_origin(typeinfo)
+    # Handle List-like type-constructor
+    # Eg. List[INT32] is a value type, while List[int] is an attribute type
+    if type_constructor in _LIST_CONSTRUCTORS:
+        elt_type = get_args(typeinfo)[0]
+        return is_value_type(elt_type)
+    # Handle Union and Optional type-constructors
+    if type_constructor is typing.Union:
+        # Filter out None, since typing.Optional[X] evaluates to Union[X, None]
+        args = [x for x in get_args(typeinfo) if x is not type(None)]  # noqa: E721
+        args_value_check = [is_value_type(x) for x in args]
+        if all(args_value_check):
+            # Handles cases like Optional[INT32] as well as Union[FLOAT16, FLOAT, DOUBLE]
+            return True
+        elif (len(args) == 1) and args_value_check[0] is False:
+            # Handle the case of optional attribute: eg. Optional[int]
+            # Note that we do not allow Union[int, float] for attributes.
+            return False
+        else:
+            raise ValueError(f"Unsupported type annotation '{typeinfo}'")
+    # Handle TypeVars:
+    if isinstance(typeinfo, typing.TypeVar):
+        if hasattr(typeinfo, "__bound__"):
+            bound = typeinfo.__bound__
+            return is_value_type(bound)
+    raise ValueError(f"Unsupported type annotation {typeinfo}")
+
+
+def is_attr_type(pytype: TypeAnnotationValue):
+    return is_value_type(pytype) is False
+
+
+def is_valid_type(typeinfo: TypeAnnotationValue):
+    try:
+        return is_value_type(typeinfo) in {True, False}
+    except ValueError:
+        return False
 
 
 def get_return_types(typeinfo: type | typing.Sequence[type]) -> typing.Sequence[type]:


Test patch: