Auto generate OpSchema for functions #594

justinchuby · 2023-04-05T17:15:54Z

This change adds the capability to auto generate OpSchema.

Changes

Implement the opschema property in OnnxFunction
Change attrs in IRFunction to self.attrs: list[IRAttributeValue] so that it contains type information
Include typeinfo in add_attr_parameter
Moved version_utils to _internal to use it in values

TODO

Test on all torch_lib functions

Example

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

onnxscript/values.py


 import onnx
 import onnx.defs

-from onnxscript import irbuilder, sourceinfo
+from onnxscript import irbuilder, sourceinfo, type_annotation


onnxscript/values.py

justinchuby · 2023-04-12T01:08:32Z

@gramalingam I also see AttrRef, how should it be used?

codecov · 2023-04-12T03:50:32Z

Codecov Report

Merging #594 (7207107) into main (f1591c7) will decrease coverage by 0.08%.
The diff coverage is 89.04%.

@@            Coverage Diff             @@
##             main     #594      +/-   ##
==========================================
- Coverage   74.20%   74.13%   -0.08%     
==========================================
  Files         107      107              
  Lines       11302    11401      +99     
  Branches     1177     1197      +20     
==========================================
+ Hits         8387     8452      +65     
- Misses       2592     2619      +27     
- Partials      323      330       +7

Impacted Files	Coverage Δ
onnxscript/_internal/version_utils.py	`100.00% <ø> (ø)`
onnxscript/converter.py	`81.74% <ø> (-1.20%)`	⬇️
onnxscript/type_annotation.py	`67.00% <55.17%> (-5.98%)`	⬇️
onnxscript/irbuilder.py	`77.12% <83.33%> (+0.17%)`	⬆️
onnxscript/onnx_types.py	`93.57% <98.07%> (+1.26%)`	⬆️
onnxscript/autocast.py	`92.85% <100.00%> (ø)`
...pt/function_libs/torch_aten/graph_building_test.py	`89.09% <100.00%> (ø)`
...s/function_libs/torch_aten/ops_correctness_test.py	`86.35% <100.00%> (-2.47%)`	⬇️
onnxscript/tests/functions/onnxfns1A_test.py	`85.45% <100.00%> (+0.26%)`	⬆️
onnxscript/tests/functions/onnxfns2_test.py	`94.73% <100.00%> (+0.14%)`	⬆️
... and 2 more

... and 1 file with indirect coverage changes

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

fatcat-z · 2023-04-12T15:58:30Z

onnxscript/onnx_types.py

@@ -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 = (


Could you please sort them in alphabetical order for better readness?

fatcat-z · 2023-04-12T16:07:28Z

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
  "Programming Language :: Python :: 3.11",
  "License :: OSI Approved :: Apache Software License",
 ]
-dependencies = ["numpy", "onnx", "typing_extensions"]
+dependencies = ["numpy", "onnx>=1.13", "typing_extensions"]


Why do we specify the minimum onnx version here?

titaiwangms · 2023-04-12T14:23:07Z

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
  "Programming Language :: Python :: 3.11",
  "License :: OSI Approved :: Apache Software License",
 ]
-dependencies = ["numpy", "onnx", "typing_extensions"]
+dependencies = ["numpy", "onnx>=1.13", "typing_extensions"]


What is this for? I see OpSchema is no older than 14?

I am planning to bump this when 14 is released. For now I can revert I guess

titaiwangms · 2023-04-12T14:30:23Z

onnxscript/values.py

        self._param_schemas: Optional[tuple[ParamSchema, ...]] = None

    def __call__(self, *args, **kwargs):
        # FIXME(after #225): Move import to the top of the file.
        from onnxscript import evaluator  # pylint: disable=import-outside-toplevel

-        return evaluator.default().eval(self.get_schema(), args, kwargs)
+        schema = self.get_schema()
+        assert schema is not None


Could this be None for an Op? What is the case here?

When the op doesn't have a schema this will fail I think. That would be the custom ops. Here we expect schema is always not None

Oh, I forgot custom ops. How does custom op get in here? I thought onnx-script didn't support non-ONNX ops. If it's a custom op built with non onnx-script way (not compmosed of ONNX standard op), how does onnx-script process it? I thought we could only handle it in exporter side.

onnxscript/values.py

titaiwangms · 2023-04-12T16:19:55Z

onnxscript/irbuilder.py

+        fn: IRFunction,
+        varname: str,
+        attribute_type: onnx.AttributeProto.AttributeType,
+        default_value: Any,


Is there a better annotation we can use than Any?

I think so. Let me refine that

titaiwangms · 2023-04-12T16:32:47Z

onnxscript/irbuilder.py

+            proto = onnx.AttributeProto()
+            proto.name = varname
+            proto.type = attribute_type
+            fn.add_attr_parameter(IRAttributeValue(proto), has_default=False)


Do you think has_defualt can be an attribute of IRAttributeValue? I think that's more straightforward.

That's a great suggestion - I thought of this too. Will make an update

titaiwangms · 2023-04-12T16:36:05Z

onnxscript/onnx_types.py

@@ -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


Do we still need this?

Not sure. Maybe?

justinchuby · 2023-04-12T18:05:11Z

Closed in favor of #626

Auto generate OpSchema for functions

a3ec431

justinchuby added module: torchlib Related to the torch/aten function lib in development topic: api labels Apr 5, 2023

justinchuby requested review from gramalingam and titaiwangms April 5, 2023 17:16

justinchuby added 11 commits April 5, 2023 17:17

Update

0224ad2

import

fc718d7

Move version utils

7240cd4

param schemas

5febf7f

This

44cd785

Merge branch 'main' into justinchu/auto-opschema

9e6e8ba

Snapshot

49b7123

Schema!

8c7e447

Working on attrs

1e058ae

Snapshot

f859b8e

Format

b591aea

justinchuby marked this pull request as ready for review April 12, 2023 00:29

justinchuby and others added 2 commits April 11, 2023 17:29

Merge branch 'main' into justinchu/auto-opschema

1c525b0

Fix

15381e4

github-advanced-security bot found potential problems Apr 12, 2023

View reviewed changes

justinchuby added 2 commits April 12, 2023 00:49

Try this

8bf732e

Try this

2c9c90e

justinchuby requested review from BowenBao and jcwchen April 12, 2023 00:56

get_supported_input_types

544f1f7

justinchuby commented Apr 12, 2023

View reviewed changes

onnxscript/values.py Outdated Show resolved Hide resolved

Update onnxscript/values.py

4334edf

justinchuby mentioned this pull request Apr 12, 2023

All the types there is #623

Open

update attributes

7fcd2e0

justinchuby added 2 commits April 12, 2023 03:39

opschema return none

ed84d6f

static_cast_inputs

d2a35bd

justinchuby added 4 commits April 12, 2023 04:07

order

4984321

type

c1bdc51

requirement

fec676b

typing

7207107

justinchuby requested review from xiaowuhu and fatcat-z April 12, 2023 05:19

fatcat-z reviewed Apr 12, 2023

View reviewed changes

titaiwangms reviewed Apr 12, 2023

View reviewed changes

justinchuby marked this pull request as draft April 12, 2023 17:33

justinchuby closed this Apr 12, 2023

justinchuby mentioned this pull request Jun 21, 2023

Unify interface for accessing custom function and op schemas #320

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Auto generate OpSchema for functions #594

Auto generate OpSchema for functions #594

justinchuby commented Apr 5, 2023 •

edited

Loading

Check notice

justinchuby commented Apr 12, 2023

codecov bot commented Apr 12, 2023 •

edited

Loading

fatcat-z Apr 12, 2023

justinchuby Apr 12, 2023

fatcat-z Apr 12, 2023

titaiwangms Apr 12, 2023

justinchuby Apr 12, 2023

titaiwangms Apr 12, 2023

justinchuby Apr 12, 2023

titaiwangms Apr 12, 2023

titaiwangms Apr 12, 2023

justinchuby Apr 12, 2023

titaiwangms Apr 12, 2023

justinchuby Apr 12, 2023

titaiwangms Apr 12, 2023

justinchuby Apr 12, 2023

justinchuby commented Apr 12, 2023

		@@ -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

Auto generate OpSchema for functions #594

Auto generate OpSchema for functions #594

Conversation

justinchuby commented Apr 5, 2023 • edited Loading

Changes

TODO

Example

Check notice

justinchuby commented Apr 12, 2023

codecov bot commented Apr 12, 2023 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

justinchuby commented Apr 12, 2023

justinchuby commented Apr 5, 2023 •

edited

Loading

codecov bot commented Apr 12, 2023 •

edited

Loading