Auto generate OpSchema for functions | feat

This change adds the capability to auto generate `OpSchema`.

### Changes

- Implement the `opschema` property in `OnnxFunction`

ghstack-source-id: 07f9122
Pull Request resolved: #626

Author: justinchuby

onnxscript/autocast.py

@@ -86,7 +86,7 @@ def cast(x, typeinfo) -> tensor.Tensor:
     return cast_inputs(get_type_info, cast, op_schema, *args)
 
 
-def static_cast_inputs(converter, op_schema: OpSchema, *args):
+def static_cast_inputs(converter, op_schema: Optional[OpSchema], *args):
     """Used for autocast during script-translation."""
     if op_schema is None:
         return args

onnxscript/onnx_types.py

@@ -105,69 +105,105 @@ def to_type_proto(cls) -> onnx.TypeProto:
             shape = [cls.shape]  # example: "FLOAT[10]"
         return onnx.helper.make_tensor_type_proto(cls.dtype, shape)
 
+    @classmethod
+    def to_string(cls) -> str:
+        raise NotImplementedError()
+
 
 class FLOAT(TensorType, dtype=onnx.TensorProto.FLOAT):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(float)"
 
 
 class UINT8(TensorType, dtype=onnx.TensorProto.UINT8):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(uint8)"
 
 
 class INT8(TensorType, dtype=onnx.TensorProto.INT8):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(int8)"
 
 
 class UINT16(TensorType, dtype=onnx.TensorProto.UINT16):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(uint16)"
 
 
 class INT16(TensorType, dtype=onnx.TensorProto.INT16):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(int16)"
 
 
 class INT32(TensorType, dtype=onnx.TensorProto.INT32):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(int32)"
 
 
 class INT64(TensorType, dtype=onnx.TensorProto.INT64):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(int64)"
 
 
 class STRING(TensorType, dtype=onnx.TensorProto.STRING):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(string)"
 
 
 class BOOL(TensorType, dtype=onnx.TensorProto.BOOL):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(bool)"
 
 
 class FLOAT16(TensorType, dtype=onnx.TensorProto.FLOAT16):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(float16)"
 
 
 class DOUBLE(TensorType, dtype=onnx.TensorProto.DOUBLE):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(double)"
 
 
 class UINT32(TensorType, dtype=onnx.TensorProto.UINT32):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(uint32)"
 
 
 class UINT64(TensorType, dtype=onnx.TensorProto.UINT64):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(uint64)"
 
 
 class COMPLEX64(TensorType, dtype=onnx.TensorProto.COMPLEX64):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(complex64)"
 
 
 class COMPLEX128(TensorType, dtype=onnx.TensorProto.COMPLEX128):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(complex128)"
 
 
 class BFLOAT16(TensorType, dtype=onnx.TensorProto.BFLOAT16):
-    pass
+    @classmethod
+    def to_string(cls):
+        return "tensor(bfloat16)"
 
 
 def onnx_type_to_onnxscript_repr(onnx_type: onnx.TypeProto) -> str:
@@ -203,3 +239,22 @@ def onnx_type_to_onnxscript_repr(onnx_type: onnx.TypeProto) -> str:
 
 # Currently, only tensor types are supported. Need to expand support for other ONNX types.
 ONNXType = TensorType
+
+ALL_TENSOR_TYPES = (
+    BFLOAT16,
+    BOOL,
+    COMPLEX128,
+    COMPLEX64,
+    DOUBLE,
+    FLOAT,
+    FLOAT16,
+    INT16,
+    INT32,
+    INT64,
+    INT8,
+    STRING,
+    UINT16,
+    UINT32,
+    UINT64,
+    UINT8,
+)

onnxscript/type_annotation.py

@@ -7,11 +7,12 @@
 import collections
 import inspect
 import typing
+from typing import Any, TypeVar, Union
 
 import onnx
 from typing_extensions import get_args, get_origin
 
-from onnxscript.onnx_types import TensorType
+from onnxscript import onnx_types
 
 # TypeAnnotationValue represents the (value of) valid type-annotations recognized
 # by ONNX Script. TODO: Flesh out a formal definition. Currently, it supports
@@ -59,7 +60,7 @@ def _is_primitive_attr_type(typeinfo: TypeAnnotationValue) -> bool:
 
 def pytype_to_attrtype(
     pytype: TypeAnnotationValue,
-) -> onnx.AttributeProto.AttributeType:
+) -> typing.Optional[onnx.AttributeProto.AttributeType]:
     pytype = _remove_annotation(pytype)
     if pytype in _PYTYPE_TO_ATTRTYPE_MAP:
         return _PYTYPE_TO_ATTRTYPE_MAP[pytype]
@@ -74,13 +75,13 @@ def pytype_to_attrtype(
         elt_type = get_args(pytype)[0]
         if elt_type in _LISTTYPE_TO_ATTRTYPE_MAP:
             return _LISTTYPE_TO_ATTRTYPE_MAP[elt_type]
-    return onnx.AttributeProto.UNDEFINED
+    return None
 
 
 def _is_tensor_type(typeinfo: TypeAnnotationValue) -> bool:
-    if isinstance(typeinfo, TensorType):
+    if isinstance(typeinfo, onnx_types.TensorType):
         return True
-    if inspect.isclass(typeinfo) and issubclass(typeinfo, TensorType):
+    if inspect.isclass(typeinfo) and issubclass(typeinfo, onnx_types.TensorType):
         return True
     return False
 
@@ -146,3 +147,52 @@ def get_return_types(typeinfo: type | typing.Sequence[type]) -> typing.Sequence[
     if get_origin(typeinfo) is tuple:
         return get_args(typeinfo)
     return (typeinfo,)
+
+
+def _reduce_type_var_to_union(hint: typing.TypeVar):
+    """Reduce a TypeVar to a Union type on which we can use issubclass to check membership."""
+    assert isinstance(hint, TypeVar)
+
+    # If the TypeVar has a bound, use that.
+    if hint.__bound__ is not None:
+        return hint.__bound__
+
+    # If the TypeVar has no bound, use the first constraint.
+    if hint.__constraints__:
+        return Union.__getitem__(hint.__constraints__)
+
+    return Any
+
+
+def get_supported_input_types(pytype) -> list[str]:
+    """Returns a list of all supported input types for a given type annotation.
+
+    Args:
+        pytype: A type annotation.
+
+    Returns:
+        A list of all supported input types for the given type annotation.
+    """
+    supported_types: list[str] = []
+    if typing.get_origin(pytype) is Union and isinstance(typing.get_args(pytype)[0], TypeVar):
+        # Recursively unpack TypeVars inside an Optional
+        for arg in typing.get_args(pytype):
+            supported_types.extend(get_supported_input_types(arg))
+        return supported_types
+
+    if isinstance(pytype, TypeVar):
+        pytype = _reduce_type_var_to_union(pytype)
+
+    for tensor_type in onnx_types.ALL_TENSOR_TYPES:
+        if pytype is None:
+            # The same as Any
+            supported_types.append(tensor_type.to_string())
+        elif pytype == onnx_types.TensorType:
+            supported_types.append(tensor_type.to_string())
+        elif isinstance(pytype, tensor_type):
+            supported_types.append(tensor_type.to_string())
+        elif issubclass(tensor_type, pytype):
+            supported_types.append(tensor_type.to_string())
+        # TODO(justinchuby): Handle sequence types
+
+    return supported_types