feat(atenlib): use typevars for functions #261
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
[ghstack-poisoned]
justinchuby
changed the title
justinchuby
added
the
module: torchlib
Codecov Report
@@ Coverage Diff @@
## main #261 +/- ##
==========================================
+ Coverage 72.27% 72.39% +0.12%
==========================================
Files 94 95 +1
Lines 9007 9021 +14
==========================================
+ Hits 6510 6531 +21
+ Misses 2497 2490 -7
Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here. |
|
@abock I think onnxscript is not complaining about the typing annotations now so we are good (hopefully?). |
This change implements
```python
TTensor = TypeVar("TTensor", bound=_TensorType)
TFloat = TypeVar("TFloat", bound=_FloatType)
TInt = TypeVar("TInt", bound=_IntType)
TReal = TypeVar("TReal", bound=_RealType)
```
and uses them to annotate the functions. This information will be used to generate symbolic functions which adds glue nodes and dispatches to the ONNX functions.
In order for type checking to work, the generated ops need to use `TypeVar`s as well.
[ghstack-poisoned]
This change implements
```python
TTensor = TypeVar("TTensor", bound=_TensorType)
TFloat = TypeVar("TFloat", bound=_FloatType)
TInt = TypeVar("TInt", bound=_IntType)
TReal = TypeVar("TReal", bound=_RealType)
```
and uses them to annotate the functions. This information will be used to generate symbolic functions which adds glue nodes and dispatches to the ONNX functions.
In order for type checking to work, the generated ops need to use `TypeVar`s as well.
[ghstack-poisoned]
fatcat-z
reviewed
Jan 3, 2023
| # abs(Tensor self) -> Tensor | ||
|
|
||
| return op.Abs(self) | ||
|
|
||
|
|
||
| @torch_op("aten::acos") | ||
| def aten_acos(self): | ||
| def aten_acos(self: TFloat) -> TFloat: |
There was a problem hiding this comment.
In current exporter implementation, actually even the input is tensor(bool) it still can work.
If we limit the input to TFloat like this, will it break export?
There was a problem hiding this comment.
The function does not take bools because onnx Acos does not. The current plan is to use the generated symbolic functions to handle dtype conversion.
fatcat-z
reviewed
Jan 3, 2023
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
This change implements
and uses them to annotate the functions. This information will be used to generate symbolic functions which adds glue nodes and dispatches to the ONNX functions.
In order for type checking to work, the generated ops need to use
TypeVars as well.