[JIT] Add support for default args in class types #42988

SplitInfinity · 2020-08-13T05:31:30Z

Stack from ghstack:

[JIT] Add support for default args in class types #42988 [JIT] Add support for default args in class types

Summary
This commit adds support for default args in methods of TorchScript
classes. Default values are already represented in the JIT TreeView
representation of function (in the defaultValue member of the Param
subclass of TreeView); this commit adds code to parse out the default
values and pack them into the TreeView created for the class definition.
This behaviour is gated by a new argument parse_defaults that has been
added to several functions in frontend.py to make sure that this new
behaviour is triggered only when TorchScript classes are being compiled.
The same code is used to compile functions and interfaces, and default
arguments are not expected to be parsed during compilation of these
entities.

Test Plan
This commit adds a unit test to TestClassType to test this feature.

python test/test_jit.py TestClassType.test_default_args

Fixes
This commit fixes #42562.

Differential Revision: D23762013

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: e3fc9c6 Pull Request resolved: #42988

dr-ci · 2020-08-13T05:41:56Z

💊 CI failures summary and remediations

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

💚 💚 Looks good so far! There are no failures yet. 💚 💚

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 154 times.

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: 8c3e916 Pull Request resolved: #42988

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: 9af3faf Pull Request resolved: #42988

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: c7fea05 Pull Request resolved: #42988

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: 085c698 Pull Request resolved: #42988

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

torch/csrc/jit/frontend/script_type_parser.cpp

torch/jit/_script.py

gmagogsfm · 2020-08-14T21:29:09Z

test/jit/test_class_type.py

+            return a.get_int() + a.get_list()[2] + a.get_tup()[1]
+
+        def some_defaults() -> int:
+            a: ClassWithDefaultArgs = ClassWithDefaultArgs(b=[5, 6, 7])


nit:
a is both an attr name and a var name here, a bit confusing.

Noted, will change the object variable name to obj

test/jit/test_class_type.py

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: 127dd00 Pull Request resolved: #42988

eellison

didnt do full review, just comment here

torch/csrc/jit/frontend/script_type_parser.cpp

eellison

What do you think about the PR here ? https://github.com/pytorch/pytorch/pull/31344/files

torch/jit/_script.py

SplitInfinity · 2020-08-18T01:03:16Z

What do you think about the PR here ? https://github.com/pytorch/pytorch/pull/31344/files

I think it's more or less what I have here. However, one thing I don't see there is checks for mutable function defaults, which, judging by the failures on this PR, happen quite a bit.

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: 87817f1 Pull Request resolved: #42988

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: 111eb48 Pull Request resolved: #42988

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. [ghstack-poisoned]

**Summary** This commit adds support for default args in methods of TorchScript classes. Default values are already represented in the JIT TreeView representation of function (in the `defaultValue` member of the `Param` subclass of `TreeView`); this commit adds code to parse out the default values and pack them into the `TreeView` created for the class definition. This behaviour is gated by a new argument `parse_defaults` that has been added to several functions in `frontend.py` to make sure that this new behaviour is triggered only when TorchScript classes are being compiled. The same code is used to compile functions and interfaces, and default arguments are not expected to be parsed during compilation of these entities. **Test Plan** This commit adds a unit test to `TestClassType` to test this feature. `python test/test_jit.py TestClassType.test_default_args` **Fixes** This commit fixes #42562. ghstack-source-id: b0398ec Pull Request resolved: #42988

SplitInfinity · 2020-09-16T03:06:57Z

Okay, I think it's ready. Major changes since last review:

Checks for mutable defaults have been moved to the AST parsing level.
Variables cannot be used as default values due to issues with closing over them. A workaround is to initialize the arg to None and then initialize with the true default value in the function body.
Overload compilation has been migrated to use default values from the AST as well. This required adding code to copy default values from the implementation decl to the overload decl (this allowed me to undo the edits to the norm kernel).

codecov · 2020-09-16T04:27:06Z

Codecov Report

Merging #42988 into gh/splitinfinity/32/base will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@                    Coverage Diff                    @@
##           gh/splitinfinity/32/base   #42988   +/-   ##
=========================================================
  Coverage                     67.95%   67.96%           
=========================================================
  Files                           384      384           
  Lines                         49768    49785   +17     
=========================================================
+ Hits                          33820    33836   +16     
- Misses                        15948    15949    +1

Impacted Files	Coverage Δ
torch/jit/_script.py	`91.33% <ø> (ø)`
.../testing/_internal/distributed/rpc/jit/rpc_test.py	`29.65% <ø> (ø)`
...g/_internal/distributed/rpc/jit/rpc_test_faulty.py	`30.43% <ø> (ø)`
torch/distributed/nn/jit/instantiator.py	`95.89% <100.00%> (+1.52%)`	⬆️
torch/jit/_recursive.py	`94.30% <100.00%> (-0.02%)`	⬇️
torch/jit/frontend.py	`91.00% <100.00%> (+0.09%)`	⬆️
torch/nn/functional.py	`92.17% <100.00%> (ø)`
torch/testing/_internal/expecttest.py	`77.55% <0.00%> (-1.03%)`	⬇️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update f5d231d...e522fb2. Read the comment docs.

SplitInfinity · 2020-09-21T16:56:34Z

Ping @eellison and @gmagogsfm

There are a lot of internal failures but I sampled a handful and they seem to be ones where the argument is annotated as float but initialized as int (or something similar). They should be easy to fix, but I'd like to deal with those that after all of the common changes look okay so that I don't have to bounce back and forth between fbcode and Github.

eellison · 2020-09-21T16:59:02Z

@SplitInfinity sure, i'll take a look today. It's a big PR... If it's not too much difficulty, it's probably worth adding special handling so that we can take an int default for a float argument.

SplitInfinity · 2020-09-21T17:13:33Z

int default value annotated float is easy but there are some other ones like int annotated as <custom type>

eellison · 2020-09-21T17:22:06Z

What's an example of the custom type one ? Yea probably just int as float is fine.

eellison

I didn't do an in-depth review yet, wanted to sync up on some aspects of the approach.

So, as far as I understand, the reason we're switching a lot of the mechanics of how we get and generate default values is because we no longer have access to an object instance of a class, as we did with nn.Modules. (please point out if this is not true).

I want to point out that this doesn't actually prohibit us from getting the default values:

closed_over_var = 1
class A(object):
    def forward(self, b=closed_over_var, c=4):
        return a  
print(A.forward.__defaults__)
# (1, 4)

This also has the nice property that we don't have to regress the scriptability of cases above with closed_over_var. The default value is bound at initialization time so there is no (semantic) reason why we couldn't support this case.

I think using the python value would help with the complexity/brittleness of this stuff

      .value("ListLiteralKind", TK_LIST_LITERAL)
      .value("DictLiteralKind", TK_DICT_LITERAL)
      .value("TupleLiteralKind", TK_TUPLE_LITERAL)
      .value("VarKind", TK_VAR)

Which atm wouldn't handle stuff like a default value'd enum, or other future things that get added.

One other question I had was - we do all of the stitching/evaluation of python defaults in the script_init.cpp file. Could we add a check for mutable values there ? Since this is only the python compilation path it shouldn't get affected by the broadcasting list defs in native_functions.yaml

I understand this is a difficult PR to land that is a really nice improvement in functionality. We can consider what aspects we consider necessary for fixing and which would just be nice to have, and then decide what to land with after.

eellison · 2020-09-17T21:52:29Z

test/distributed/nn/jit/test_instantiator.py

        ) = instantiator.get_arg_return_types_from_interface(MyModuleInterface)
        self.assertEqual(args_str, "tensor, number, word")
-        self.assertEqual(arg_types_str, "tensor: Tensor, number: int, word: str")
+        self.assertEqual(arg_types_str, "tensor: Tensor, number: int, word: str = \"default\"")


why was this change needed ?

I think this is some test that checks that a schema is parsed in a specific way and the test was written in a way that didn't include defaults (because they weren't being parsed).

eellison · 2020-09-17T21:52:57Z

test/test_jit.py

-
        @torch.jit.script
-        def bool_fn(x, a=outer_c, flag=outer_flag):
+        def bool_fn(x, a: Tensor = torch.tensor(9), flag: Tensor = torch.tensor(False)):


Why was this change needed ?

This was needed due to issue with closing over variables used as default argument values.

eellison · 2020-09-17T21:53:20Z

test/test_jit.py

        """
        @torch.jit.script
-        def foo(x=torch.ones(1)):
+        def foo(x: Tensor = torch.ones(1)):


Why was this change needed ?

Now that the defaults are being parsed from the AST, there's no way to infer types for the default argument expressions without making a small graph containing them and running it.

SplitInfinity · 2020-09-21T20:42:57Z

I didn't do an in-depth review yet, wanted to sync up on some aspects of the approach.

So, as far as I understand, the reason we're switching a lot of the mechanics of how we get and generate default values is because we no longer have access to an object instance of a class, as we did with nn.Modules. (please point out if this is not true).

I want to point out that this doesn't actually prohibit us from getting the default values:
closed_over_var = 1
class A(object):
    def forward(self, b=closed_over_var, c=4):
        return a  
print(A.forward.__defaults__)
# (1, 4)

Yeah, I think we could pack all those defaults into a some map and pass it into jit_script_class_compile and then reuse getSchemaWithNameAndDefaults on all of the methods and properties and their defaults contained in that map. The only disadvantage to this approach is that the default values are not known at compile time because they are added to the schema after compilation. But that's the case with scripted functions anyway. I wish I had thought of that first T_T

One other question I had was - we do all of the stitching/evaluation of python defaults in the script_init.cpp file. Could we add a check for mutable values there ? Since this is only the python compilation path it shouldn't get affected by the broadcasting list defs in native_functions.yaml

I think we could, but we would do the check after compilation. By doing the check in frontend.py, we are able to check for that and avoid compilation if any defaults are mutable.

eellison · 2020-09-22T00:40:31Z

I think we could, but we would do the check after compilation. By doing the check in frontend.py, we are able to check for that and avoid compilation if any defaults are mutable.

I don't think this is a compelling reason for the extra complexity and maintenance. How often do people use default values that we're optimizing for their failure case? It would be just as easy to iterate over the default values before compilation in the other approach.

The only disadvantage to this approach is that the default values are not known at compile time because they are added to the schema after compilation. But that's the case with scripted functions anyway. I wish I had thought of that first T_T

Yea, it's true that this PR does add that functionality, which is really nice. It does regress other aspects of scripting though (closing over named constants and int as floats is the big one). I would prefer if we tried to improve that aspect of scripting in a separate PR.

SplitInfinity · 2020-09-22T00:50:56Z

I just posted #45098 that addresses the same GitHub issue as this PR but with the same approach that we use now for methods of modules and functions (getting the Python objects for all of the default args, converting them to IValues and attaching them to the schemas of the class methods). Maybe it's better to land that instead of this one?

SplitInfinity · 2020-09-24T19:57:00Z

F

SplitInfinity requested a review from apaszke as a code owner August 13, 2020 05:31

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

SplitInfinity requested review from mrshenli, pietern and zhaojuanmao as code owners August 14, 2020 03:43

SplitInfinity requested review from eellison and gmagogsfm August 14, 2020 20:51

gmagogsfm reviewed Aug 14, 2020

View reviewed changes

eellison reviewed Aug 17, 2020

View reviewed changes

torch/csrc/jit/frontend/script_type_parser.cpp Outdated Show resolved Hide resolved

eellison reviewed Aug 18, 2020

View reviewed changes

torch/jit/_script.py Outdated Show resolved Hide resolved

SplitInfinity requested review from eellison and gmagogsfm September 16, 2020 03:07

eellison reviewed Sep 21, 2020

View reviewed changes

SplitInfinity mentioned this pull request Sep 22, 2020

[JIT] Add default arguments for class types #45098

Closed

SplitInfinity closed this Sep 24, 2020

facebook-github-bot deleted the gh/splitinfinity/32/head branch October 25, 2020 14:18

[JIT] Add support for default args in class types #42988

[JIT] Add support for default args in class types #42988

Uh oh!

Conversation

SplitInfinity commented Aug 13, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

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

Uh oh!

💊 CI failures summary and remediations

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!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

SplitInfinity commented Aug 18, 2020

Uh oh!

SplitInfinity commented Sep 16, 2020

Uh oh!

codecov bot commented Sep 16, 2020 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

SplitInfinity commented Sep 21, 2020

Uh oh!

eellison commented Sep 21, 2020

Uh oh!

SplitInfinity commented Sep 21, 2020

Uh oh!

eellison commented Sep 21, 2020

Uh oh!

eellison left a comment

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!

SplitInfinity commented Sep 21, 2020

Uh oh!

eellison commented Sep 22, 2020

Uh oh!

SplitInfinity commented Sep 22, 2020

Uh oh!

SplitInfinity commented Sep 24, 2020

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

SplitInfinity commented Aug 13, 2020 •

edited

Loading

dr-ci bot commented Aug 13, 2020 •

edited

Loading

codecov bot commented Sep 16, 2020 •

edited

Loading