-
Notifications
You must be signed in to change notification settings - Fork 26.3k
[jit] Resolve NamedTuple types in Python #26443
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.
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
Conversation
eellison
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Cool! Thanks for doing this.
lgtm if tests pass. as a follow up we can try this with generic UDT besides NamedTuples - seems like the logic would work as well.
| elif hasattr(ann, "__torch_script_class__"): | ||
| return ClassType(_qualified_name(ann)) | ||
| elif lookup is not None: | ||
| # Maybe resolve a NamedTuple to a Tuple Type |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It's okay if this is limited to NamedTuple types for now, however I think the logic extends to any custom type...
torch/jit/annotations.py
Outdated
|
|
||
|
|
||
| def ann_to_type(ann): | ||
| def ann_to_type(ann, lookup=None): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe more descriptive name than lookup ? or add comment that it's an (rcb, sourcerange) tuple
|
i think tests didn't run, rebase ? |
Summary: When used as annotations on Python functions, `NamedTuple`s go through our Python annotation -> type mapping which previously had no way of lookup up `NamedTuple`s (which are created lazily by checking if the type has certain properties, so the lookup is creating the `TupleType` from scratch). This PR threads through the necessary data to make them work. Fixes pytorch#26437 Pull Request resolved: pytorch#26443 Pulled By: driazati Differential Revision: D17486441 fbshipit-source-id: a6bbb543ff05a5abe61f1a7f68db9ecdb652b358
When used as annotations on Python functions,
NamedTuples go through our Python annotation -> type mapping which previously had no way of lookup upNamedTuples (which are created lazily by checking if the type has certain properties, so the lookup is creating theTupleTypefrom scratch). This PR threads through the necessary data to make them work.Fixes #26437
Differential Revision: D17486441