[1/N] Implement Enum JIT support #41390

gmagogsfm · 2020-07-14T08:27:14Z

Add EnumType and AnyEnumType as first-class jit type
Add Enum-typed IValue
Enhanced aten::eq to support Enum

Supported:
Enum-typed function targuments
using Enum type and comparing them

TODO:
Add PyThon sugared value for Enum
Support getting name/value attrs of enums
Support Enum-typed return values
Support enum values of different types in same Enum class
Support serialization and deserialization

dr-ci · 2020-07-14T08:30:36Z

💊 CI failures summary and remediations

As of commit c642907 (more details on the Dr. CI page):

1/1 failures possibly* introduced in this PR
- 1/1 non-CircleCI failure(s)

ci.pytorch.org: 1 failed

Failed: pr/pytorch-linux-xenial-rocm3.5.1-py3.6

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 127 times.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

eellison

Looks like a great start!! Just a few comments here or there. Overall looks great though. LMK if you need any pointers for serialization / attribute access.

aten/src/ATen/core/ivalue.cpp

aten/src/ATen/core/ivalue_inl.h

eellison · 2020-07-15T21:23:31Z

aten/src/ATen/core/jit_type.h

Didn't think of this earlier, but I don't think we should support mutable Enum values. I don't think it's a common use case, and we would have to maintain one singleton IValue.

class Color(Enum): RED = torch.tensor(1) # all Color.RED instances would have to be updated, or we have to maintain a single instance Color.RED += 2

For the same reason, I don't know if we should support Any because it could be a mutable type (and it's also not a common use case).

Yeah, current implementation would break with mutable values. Removed Tensor and Any type.

aten/src/ATen/core/jit_type.h

eellison · 2020-07-15T21:26:36Z

aten/src/ATen/core/jit_type.h

Can we force the QualifiedName to always be set ?

Good point. Added an assert to enforce that.

eellison · 2020-07-15T21:27:43Z

torch/_jit_internal.py

Why did we need to make this change ? Seems a little strange

Enum class doesn't have "__name__" attr, instead they use "name", I am not sure why the inconsistency.

Weird. nit: maybe elif isinstance(obj, Enum) in elif conditional instead so it's a little more clear why this is needed

Indeed very weird. Done.

eellison · 2020-07-15T21:30:09Z

torch/jit/annotations.py

see comment above about supporting mutable types for enum value

So AnyType is no longer appropriate. Should I create another type that is super-type of int, float and string for this case?

I think we should just throw an error instead. If users complain we can get back to this, but the common case is one type of value for an enum.

Good point. Done.

torch/jit/annotations.py

aten/src/ATen/core/jit_type.h

gmagogsfm

Thanks for the suggestions, I have modified accordingly. Please take another look.

aten/src/ATen/core/jit_type.h

gmagogsfm · 2020-07-15T22:33:15Z

aten/src/ATen/core/jit_type.h

Yeah, current implementation would break with mutable values. Removed Tensor and Any type.

gmagogsfm · 2020-07-15T23:09:00Z

aten/src/ATen/core/jit_type.h

Good point. Added an assert to enforce that.

gmagogsfm · 2020-07-15T23:13:41Z

torch/_jit_internal.py

Enum class doesn't have "__name__" attr, instead they use "name", I am not sure why the inconsistency.

torch/jit/annotations.py

gmagogsfm · 2020-07-15T23:23:34Z

torch/jit/annotations.py

So AnyType is no longer appropriate. Should I create another type that is super-type of int, float and string for this case?

aten/src/ATen/core/jit_type.h

aten/src/ATen/core/ivalue.cpp

aten/src/ATen/core/ivalue_inl.h

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

eellison

Looks good as far as logic in this PR. I think we should wait to land this until the rest of the stack is implemented, bc otherwise someone could try to use Enums and get half-implemented version (as we have sometimes done with other features :C)

aten/src/ATen/core/ivalue_inl.h

eellison · 2020-07-16T22:56:30Z

torch/_jit_internal.py

Weird. nit: maybe elif isinstance(obj, Enum) in elif conditional instead so it's a little more clear why this is needed

eellison · 2020-07-16T22:57:42Z

torch/jit/annotations.py

I think we should just throw an error instead. If users complain we can get back to this, but the common case is one type of value for an enum.

SplitInfinity

Some initial comments

SplitInfinity · 2020-07-16T22:30:05Z

aten/src/ATen/core/ivalue.h

It seems like there are formatting changes included with your addition below. I recommend undoing the formatting changes.

Yeah, I actually intentionally added them. These are all at or near codes that I touched in this PR, in the spirit of leaving the camp ground cleaner. Let me know if this complies with project convention or not.

I personally prefer to make separate commits to clean up formatting but it's not a big deal.

Ack. Thanks.

SplitInfinity · 2020-07-16T22:30:34Z

aten/src/ATen/core/ivalue.h

Nit: undo formatting

SplitInfinity · 2020-07-16T22:30:38Z

aten/src/ATen/core/ivalue.h

Nit: undo formatting

SplitInfinity · 2020-07-16T22:35:31Z

aten/src/ATen/core/ivalue_inl.h

Was this intentionally added?

Yeah, I noticed some of these functions are separated by empty lines but some are not, so made them more consistent.

SplitInfinity · 2020-07-16T22:35:53Z

aten/src/ATen/core/ivalue_inl.h

Was this intentionally added?

SplitInfinity · 2020-07-16T22:37:01Z

aten/src/ATen/core/ivalue_inl.h

Consider storing this as a c10::QualifiedName. That class has useful methods for extracting the prefix, unqualified name, etc. that might be useful later.

Ah, nice! Done.

SplitInfinity · 2020-07-16T22:40:33Z

aten/src/ATen/core/jit_type.h

Why is the name allowed to be optional? All enum types must have names, right?

Also weren't there plans to support heterogeneous enums with mixed value types? Or is that coming later or out of scope?

Name isn't allowed to be empty. "optional" was originally added to conform to NamedType's convention, but now I realized constructor can take concrete value and then create "optional" to make NamedType happy. Changed.

To answer your second question, Elias gave feedback that heterogeneous enum value types is quite uncommon and makes enum logic a lot more complicated, so we decided to mark it out of scope for now.

SplitInfinity · 2020-07-16T22:53:58Z

test/jit/test_enum.py

Perhaps this is coming in the future but you should consider adding a test that checks that enums can be returned from scripted functions (i.e. copied back correctly from TS to Python).

Definitely, this is listed as a TODO in the PR description.

SplitInfinity · 2020-07-16T22:57:35Z

aten/src/ATen/core/jit_type.h

Suggested change

value_type_(value_type), cu_(cu) {

value_type_(std::move(value_type)), cu_(cu) {

Done. But out of curiosity, value_type is already a light-weight TypePtr, does it still benefit from std::move?

TypePtr is defined as using TypePtr = std::shared_ptr<Type> so I think std::move is a good idea here. Copying it requires atomically incrementing the reference count.

gmagogsfm · 2020-07-17T04:49:32Z

Looks good as far as logic in this PR. I think we should wait to land this until the rest of the stack is implemented, bc otherwise someone could try to use Enums and get half-implemented version (as we have sometimes done with other features :C)

I see your point, but I am a bit concerned about the tsunami of merge conflicts because these PRs combined will touch a lot of files.

How about this? I can hide this functionality behind say a environment variable-based config knob, so that no users would "accidentally" access this feature, and we can avoid merge conflicts.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

eellison

I think there'll probably be less merge conflicts than you think, but sure sounds good to land.
Could you add two tests to make sure this is properly guarded:
Something like this fails:

 def enum_comp(x: Color, y: Color) -> bool:
            return x == y
torch.jit.script(enum_comp)

and something like this fails:

class Mod(Module):
      def __init__(self):
            ...
            self.enum_attribute = Color....

      def forward(self, x):
              return self.enum_attribute
torch.jit.script(Mod())

eellison · 2020-07-17T17:08:08Z

torch/csrc/jit/serialization/unpickler.cpp

AnyEnumType will never show up as an instantiated type, it can go with AnyTupleType::Kind in set of cases

eellison · 2020-07-17T17:19:36Z

torch/jit/annotations.py

It'd be better to use https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/core/jit_type.h#L1513 so that we get a good error message here ("Type str did not match type int...") . you'd need to convert the types to jit type first with ann_to_type.

eellison · 2020-07-17T17:19:53Z

torch/jit/annotations.py

Return None here to disable

Added a check here, but raised an exception instead of returning None so that user gets a clearer error.

gmagogsfm

Added a feature guard to prevent users from accidentally using it.

gmagogsfm · 2020-07-17T18:10:28Z

aten/src/ATen/core/ivalue.h

Ack. Thanks.

gmagogsfm · 2020-07-17T18:10:52Z

torch/csrc/jit/serialization/unpickler.cpp

gmagogsfm · 2020-07-17T20:24:18Z

torch/jit/annotations.py

gmagogsfm · 2020-07-17T20:24:55Z

torch/jit/annotations.py

Added a check here, but raised an exception instead of returning None so that user gets a clearer error.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

eellison · 2020-07-17T22:55:23Z

@gmagogsfm not sure if you requested review intentionally, but generally the workflow on pytorch is a reviewer will accept a PR with small outstanding fixes if they think they're easily addressable / nits that dont necessarily have to be addressed.

Can take another look if you'd like but you're good to land without me looking

gmagogsfm · 2020-07-17T22:58:28Z

@gmagogsfm not sure if you requested review intentionally, but generally the workflow on pytorch is a reviewer will accept a PR with small outstanding fixes if they think they're easily addressable / nits that dont necessarily have to be addressed.

Can take another look if you'd like but you're good to land without me looking

Yeah, I clicked requested review to be on the safe side. But now that I think about it, the additional changes are tiny, so I will and now. Thanks eelison and splitinfinity for reviewing.

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

* Enum support tempoartily hidden behind environemnt variable EXPERIMENTAL_ENUM_SUPPORT to avoid misuse * Add EnumType and AnyEnumType as first-class jit type * Add Enum-typed IValue * Enhanced aten::eq to support Enum Supported: Enum-typed function arguments using Enum type and comparing them TODO: Add PyThon sugared value for Enum Support getting name/value attrs of enums Support Enum-typed return values Support enum values of different types in same Enum class Support serialization and deserialization

facebook-github-bot

@gmagogsfm has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot · 2020-07-19T06:15:35Z

@gmagogsfm merged this pull request in 4a3aad3.

facebook-github-bot added the oncall: jit Add this issue/PR to JIT oncall triage queue label Jul 14, 2020

facebook-github-bot reviewed Jul 14, 2020

View reviewed changes

gmagogsfm marked this pull request as ready for review July 15, 2020 00:26

gmagogsfm requested a review from apaszke as a code owner July 15, 2020 00:26

gmagogsfm requested a review from eellison July 15, 2020 00:27

eellison reviewed Jul 15, 2020

View reviewed changes

gmagogsfm commented Jul 16, 2020

View reviewed changes

gmagogsfm requested a review from eellison July 16, 2020 00:37

facebook-github-bot reviewed Jul 16, 2020

View reviewed changes

eellison reviewed Jul 16, 2020

View reviewed changes

SplitInfinity reviewed Jul 16, 2020

View reviewed changes

facebook-github-bot reviewed Jul 17, 2020

View reviewed changes

eellison approved these changes Jul 17, 2020

View reviewed changes

gmagogsfm commented Jul 17, 2020

View reviewed changes

gmagogsfm requested a review from eellison July 17, 2020 20:30

facebook-github-bot reviewed Jul 17, 2020

View reviewed changes

facebook-github-bot reviewed Jul 18, 2020

View reviewed changes

facebook-github-bot closed this in 4a3aad3 Jul 19, 2020

facebook-github-bot added the merged label Jul 19, 2020

gmagogsfm linked an issue Jul 20, 2020 that may be closed by this pull request

Support Enum in TorchScript #41095

Closed

mruberry added the Merged label Oct 28, 2020

	value_type_(value_type), cu_(cu) {
	value_type_(std::move(value_type)), cu_(cu) {

[1/N] Implement Enum JIT support #41390

[1/N] Implement Enum JIT support #41390

Uh oh!

Conversation

gmagogsfm commented Jul 14, 2020

Uh oh!

dr-ci bot commented Jul 14, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

💊 CI failures summary and remediations

ci.pytorch.org: 1 failed

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

gmagogsfm left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

facebook-github-bot left a comment

Choose a reason for hiding this comment

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

dr-ci bot commented Jul 14, 2020 •

edited

Loading