Skip to content

Create the OpLike protocol and refactor Op | feat(values) #692

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 10 commits into from
Closed

Create the OpLike protocol and refactor Op | feat(values) #692

wants to merge 10 commits into from

Conversation

justinchuby
Copy link
Collaborator

@justinchuby justinchuby commented Apr 27, 2023 •
edited
Loading

Stack from ghstack (oldest at bottom):

Changes

  • Creates the OpLike protocol that defines common attributes and methods for Op, OnnxFunction and TraceOnlyFunction so we can assume a common interface.
  • Implement param_schemas for TraceOnlyFunction
  • Refactor param_schemas to extract common logic.
  • Removes is_single_op from Op because it is unused.

Signed-off-by: Justin Chu [email protected]

@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
3dbb59e 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <[email protected]>

[ghstack-poisoned]
This was referenced Apr 27, 2023
Merged
Merged
Merged
Merged
Merged
justinchuby added a commit that referenced this pull request Apr 27, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
4c3dea1 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 2df23c2
Pull Request resolved: #692
@justinchuby justinchuby added module: torchlib Related to the torch/aten function lib in development topic: api labels Apr 27, 2023
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
a2dbc81 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 27, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
715bacb 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 7442033
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
justinchuby added a commit that referenced this pull request Apr 27, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
12fe612 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 1f78a26
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
gramalingam
gramalingam reviewed Apr 27, 2023
View reviewed changes
onnxscript/function_libs/torch_lib/tracing.py
@@ -33,7 +37,7 @@ def _get_src_and_ast(f: types.FunctionType) -> tuple[str, ast.FunctionDef]:
return src, f_ast


class TraceOnlyFunction:
class TraceOnlyFunction(onnxscript.values.OpLike):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What are the differences between a TraceOnlyFunction and an OnnxFunction? Eg., would it make sense for an OnnxFunction to be an extension (or derived class) of a TraceOnlyFunction?

Copy link
Collaborator Author

@justinchuby justinchuby Apr 27, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A TraceOnlyFunction is simply a python function with added attributes to make it more like an onnx function. To me it is a semantically different thing that happens to implement the same protocol for convenience. But I could be wrong.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let me see if a TraceOnlyFunction can be derived from Op

titaiwangms
titaiwangms reviewed Apr 27, 2023
View reviewed changes
Copy link
Contributor

@titaiwangms titaiwangms left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I suggest we move TraceOnlyFunction to values.py. It would make more sense to the existence of OpLike. And if we decide to go with the current setting, we should change the file name tracing to something more specific. Tracing seems general... Codes are all good!

@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
ae2d836 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 27, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
ac4aa6e 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 1f78a26
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
justinchuby added a commit that referenced this pull request Apr 27, 2023
@justinchuby
Implement split_function_and_wrangler | test(torchlib) (#688)
a05ab66 
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at
bottom):
* #692
* #674
* #684
* #626
* #631
* __->__ #688
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
a5e9975 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
fbfbc29 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 1f78a26
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
561b2b7 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
c5a15cc 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 1f78a26
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
b88a8ed 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
357e18c 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: bf354dd
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Analyze type annotation for input types | feat(op_schema) (#631)
0b633ea 
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at
bottom):
* #692
* #674
* #626
* #684
* __->__ #631
* #688
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
aa1e4a7 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
b18eb58 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: ff93f5d
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Refactor converter to isolate translate_function_signature logic | fe…
d3ce597 
…at(converter) (#684)

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at
bottom):
* #692
* #674
* #626
* __->__ #684

This change refactors the `translate_function_def` method in `Converter`
to isolate the signature handling logic to
`translate_function_signature`. `translate_function_signature` is used
in #674 to handle function signatures so we do not need to translate the
function body for general python functions incompatible with ONNX.

Signed-off-by: Justin Chu <[email protected]>
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
fdc1327 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
575f5e7 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: c47cc99
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
b811484 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
962fe56 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: c47cc99
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Auto generate OpSchema for functions | feat(op_schema) (#626)
2a73b15 
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at
bottom):
* #692
* #674
* __->__ #626
* #684

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
@justinchuby
Update on "Create the OpLike protocol and refactor Op | feat(values)"
b94d11b 
## Changes

- Creates the `OpLike` protocol that defines common attributes and methods for `Op`, `OnnxFunction` and `TraceOnlyFunction` so we can assume a common interface.
- Implement `param_schemas` for `TraceOnlyFunction`
- Refactor `param_schemas` to extract common logic.
- Removes `is_single_op` from `Op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

[ghstack-poisoned]
@justinchuby justinchuby mentioned this pull request Apr 28, 2023
Merged
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
39181b7 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 853c076
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
justinchuby added a commit that referenced this pull request Apr 28, 2023
@justinchuby
Create the OpLike protocol and refactor Op | feat(values)
02edd76 
- Removes `is_single_op` because it is unused.

Signed-off-by: Justin Chu <justinchumicrosoft.com>

ghstack-source-id: 853c076
Pull Request resolved: #692

Signed-off-by: Justin Chu <[email protected]>
@justinchuby justinchuby deleted the gh/justinchuby/23/head branch March 13, 2024 01:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@gramalingam gramalingam gramalingam left review comments

@titaiwangms titaiwangms titaiwangms approved these changes

@BowenBao BowenBao Awaiting requested review from BowenBao

Assignees
No one assigned
Labels
module: torchlib Related to the torch/aten function lib in development topic: api
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@justinchuby @gramalingam @titaiwangms