Update on "Auto generate OpSchema for functions | feat(op_schema)"

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

### Changes

- Implement the `opschema` property in `OnnxFunction`
- Test on all torch_lib functions

### Next PR

Support trace_only functions

## Example

```python
from onnxscript.function_libs.torch_aten.ops import core, nn


print("core.aten_abs.opschema: ", core.aten_abs.opschema)

print("nn.aten_cross_entropy_loss.opschema: ", nn.aten_cross_entropy_loss.opschema)
```

Results

```
core.aten_abs.opschema:  OpSchema(
    name='aten_abs',
    domain='onnxscript.atenlib',
    since_version=1,
    doc='abs(Tensor self) -> Tensor',
    type_constraints=[OpSchema.TypeConstraintParam(type_param_str='TReal', allowed_type_strs=['tensor(float)', 'tensor(int8)', 'tensor(int16)', 'tensor(int32)', 'tensor(int64)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description='')],
    inputs=[OpSchema.FormalParameter(name='self', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    outputs=[OpSchema.FormalParameter(name='return_val', type_str='TReal', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    attributes={}
)
nn.aten_cross_entropy_loss.opschema:  OpSchema(
    name='aten_cross_entropy_loss',
    domain='onnxscript.atenlib',
    since_version=1,
    doc='cross_entropy_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, float label_smoothing=0.0) -> Tensor',
    type_constraints=[OpSchema.TypeConstraintParam(type_param_str='TFloatOrBFloat16', allowed_type_strs=['tensor(float)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description=''), OpSchema.TypeConstraintParam(type_param_str='T1', allowed_type_strs=['tensor(float)', 'tensor(float16)', 'tensor(double)', 'tensor(bfloat16)'], description='')],
    inputs=[OpSchema.FormalParameter(name='self', type_str='TFloatOrBFloat16', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>), OpSchema.FormalParameter(name='weight', type_str='T1', description='', param_option=<FormalParameterOption.Optional: 1>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    outputs=[OpSchema.FormalParameter(name='result_10', type_str='TFloatOrBFloat16', description='', param_option=<FormalParameterOption.Single: 0>, is_homogeneous=True, min_arity=1, differentiation_category=<DifferentiationCategory.Unknown: 0>)],
    attributes={'ignore_index': OpSchema.Attribute(name='ignore_index', type=<AttrType.INT: 2>, description='', default_value=name: "ignore_index"
i: -100
type: INT
, required=False), 'label_smoothing': OpSchema.Attribute(name='label_smoothing', type=<AttrType.FLOAT: 1>, description='', default_value=name: "label_smoothing"
f: 0.0
type: FLOAT
, required=False), 'reduction': OpSchema.Attribute(name='reduction', type=<AttrType.INT: 2>, description='', default_value=name: "reduction"
i: 1
type: INT
, required=False), 'target': OpSchema.Attribute(name='target', type=<AttrType.INTS: 7>, description='', default_value=, required=True)}
)
```

Fixes #476

[ghstack-poisoned]

Author: justinchuby

onnxscript/type_annotation.py

@@ -44,6 +44,23 @@
 
 _LIST_CONSTRUCTORS = frozenset([list, typing.List, typing.Sequence, collections.abc.Sequence])
 
+ALL_TYPE_STRINGS = (
+    "tensor(bfloat16)",
+    "tensor(bool)",
+    "tensor(double)",
+    "tensor(float)",
+    "tensor(float16)",
+    "tensor(int16)",
+    "tensor(int32)",
+    "tensor(int64)",
+    "tensor(int8)",
+    "tensor(string)",
+    "tensor(uint16)",
+    "tensor(uint32)",
+    "tensor(uint64)",
+    "tensor(uint8)",
+)
+
 
 def _remove_annotation(typeinfo: TypeAnnotationValue) -> TypeAnnotationValue:
     """Remove Annotated wrapper if present, otherwise return typeinfo as is."""

onnxscript/type_annotation_test.py

@@ -3,20 +3,17 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
-from typing import Optional, Sequence, TypeVar, Union
 import unittest
+from typing import Any, Optional, Sequence, TypeVar, Union
 
 import parameterized
 
-import onnxscript.testing
 import onnxscript
-from onnxscript import script
+import onnxscript.testing
+from onnxscript import FLOAT, INT64, script, type_annotation
 from onnxscript.onnx_opset import opset15 as op
-from onnxscript import FLOAT, INT64
 from onnxscript.tests.common import testutils
 
-from onnxscript import type_annotation
-
 
 class TypeAnnotationTest(testutils.TestBase):
     def test_type_annotation(self):
@@ -99,7 +96,7 @@ def bool_type_for_attribute(self: FLOAT[...], sorted: bool) -> FLOAT[...]:
 _TestTypeVarTwoBound = TypeVar("_TestTypeVarTwoBound", bound=Union[INT64, FLOAT])
 
 
-class UtilityFunctionsTest(unittest.TestCase):
+class TypeConversionFunctionsTest(unittest.TestCase):
     @parameterized.parameterized.expand(
         [
             (
@@ -121,10 +118,12 @@ class UtilityFunctionsTest(unittest.TestCase):
             (
                 "optional_tensor_type_all",
                 Optional[onnxscript.onnx_types.TensorType],
-                type_annotation.ALL_TYPE_STRINGS
-                + [
-                    f"optional({tensor_type})"
-                    for tensor_type in type_annotation.ALL_TYPE_STRINGS
+                [
+                    *type_annotation.ALL_TYPE_STRINGS,
+                    *[
+                        f"optional({tensor_type})"
+                        for tensor_type in type_annotation.ALL_TYPE_STRINGS
+                    ],
                 ],
             ),
             (
@@ -214,7 +213,7 @@ class UtilityFunctionsTest(unittest.TestCase):
             ),
         ]
     )
-    def test_pytype_to_input_strings(self, _, pytype: Any, expected):
+    def test_pytype_to_input_strings(self, _, pytype: Any, expected: list[str]):
         self.assertEqual(type_annotation.pytype_to_input_strings(pytype), expected)
 
     @parameterized.parameterized.expand(
@@ -231,15 +230,15 @@ def test_pytype_to_input_strings(self, _, pytype: Any, expected):
                 Sequence[_TestTypeVarOneBound],
                 "Sequence_TestTypeVarOneBound",
             ),
-            ("normal_type", INT64, "None"),
+            ("normal_type", INT64, None),
             ("union_type", Union[INT64, FLOAT], None),
             ("optional_type", Optional[INT64], None),
             ("sequence_type", Sequence[INT64], None),
             ("optional_sequence_type", Optional[Sequence[INT64]], None),
             ("optional_union_type", Optional[Union[INT64, FLOAT]], None),
         ]
     )
-    def get_type_constraint_name(self, _: str, pytype, expected):
+    def get_type_constraint_name(self, _: str, pytype: Any, expected: Optional[str]):
         self.assertEqual(type_annotation.get_type_constraint_name(pytype), expected)