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_test.py

@@ -3,14 +3,20 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
+from typing import TypeVar, Union
 import unittest
 
+import parameterized
+
 import onnxscript.testing
+import onnxscript
 from onnxscript import script
 from onnxscript.onnx_opset import opset15 as op
 from onnxscript.onnx_types import FLOAT
 from onnxscript.tests.common import testutils
 
+from onnxscript import type_annotation
+
 
 class TypeAnnotationTest(testutils.TestBase):
     def test_type_annotation(self):
@@ -87,9 +93,27 @@ def bool_type_for_attribute(self: FLOAT[...], sorted: bool) -> FLOAT[...]:
             bool_type_for_attribute, bool_type_for_attribute_txt
         )
 
+_TestTypeVarConstraints = TypeVar("_TestTypeVarConstraints", onnxscript.INT64, onnxscript.FLOAT)
+_TestTypeVarOneBound = TypeVar("_TestTypeVarOneBound", bound=onnxscript.INT64)
+_TestTypeVarTwoBound = TypeVar("_TestTypeVarTwoBound", bound=Union[onnxscript.INT64, onnxscript.FLOAT])
+
 
 class UtilityFunctionsTest(unittest.TestCase):
-    def test_pytype_to_input_strings(self):
+    @parameterized.parameterized.expand(
+        [
+            ("tensor_type", onnxscript.onnx_types.TensorType, type_annotation.ALL_TYPE_STRINGS),
+            ("tensor_type", onnxscript.INT64, ["tensor(int64)"]),
+            ("tensor_type_variadic_shape", onnxscript.INT64[...], ["tensor(int64)"]),
+            ("tensor_type_shape", onnxscript.INT64[10], ["tensor(int64)"]),
+            ("type_var_constraints", _TestTypeVarConstraints, ["tensor(int64)"]),
+            ("type_bound_one", _TestTypeVarOneBound, ["tensor(int64)"]),
+            ("type_bound_two", _TestTypeVarTwoBound, ["tensor(int64)", "tensor(float)"]),
+        ]
+    )
+    def test_pytype_to_input_strings(self, _, pytype: Any, expected)
+        pass
+
+    def get_type_constraint_name(self, _: str, typevar, expected):
         pass
 
 if __name__ == "__main__":