torch.jit.script escape hatch #106229

pmeier · 2023-07-28T21:06:34Z

Although the sun is setting for torchscript, it is not officially deprecated since nothing currently fully replaces it. Thus, "downstream" libraries like TorchVision, that started offering torchscript support still need to support it for BC.

torchscript has forced us to use workaround after workaround since forever. Although this makes the code harder to read and maintain, we made our peace with it. However, we are currently looking into more elaborate API designs that are severely hampered by our torchscript BC guarantees.

Although likely not intended as such, while looking for ways to enable our design while keeping a subset of it scriptable, we found the undocumented __prepare_scriptable__ escape hatch:

pytorch/torch/jit/_script.py

Line 977 in 0cf9189

    
           obj = obj.__prepare_scriptable__() if hasattr(obj, '__prepare_scriptable__') else obj  # type: ignore[operator]

One can define this method and if you call torch.jit.script on the object, the returned object of the method will be scripted rather than the original object. In TorchVision we are using exactly this mechanism to enable BC while allowing the object in eager mode to be a lot more flexible (*args, **kwargs, dynamic dispatch, ...).

Unfortunately, this escape hatch is only available for nn.Module's

pytorch/torch/jit/_script.py

Lines 1279 to 1283 in 0cf9189

    
           if isinstance(obj, torch.nn.Module): 
        
               obj = call_prepare_scriptable_func(obj) 
        
               return torch.jit._recursive.create_script_module( 
        
                   obj, torch.jit._recursive.infer_methods_to_compile 
        
               )

This was fine for the example above since we were subclassing from nn.Module anyway. However, we recently also hit a case where this wasn't the case.

Given the frozen state on JIT, would it be possible to give us a general escape hatch so that we can move forward with the design unconstrained while still keeping BC?

This PR implements just this by re-using the __prepare_scriptable__ hook.

pytorch-bot · 2023-07-28T21:06:37Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106229

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 1fcca1a:

NEW FAILURE - The following job has failed:

linux-bionic-py3.8-clang9 / test (dynamo, 1, 2, linux.2xlarge) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

lezcano

SGTM as this unblocks vision, and this feature is not documented anyway, but let's see what @ezyang thinks.

ezyang

OK

pmeier · 2023-08-11T18:22:49Z

@pytorchbot merge -f "unrelated dynamo failure"

pytorchmergebot · 2023-08-11T18:24:40Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/) [ghstack-poisoned]

This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/) ghstack-source-id: 216795799 Pull Request resolved: #120806

…sures" This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/) [ghstack-poisoned]

This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/) [ghstack-poisoned]

…sures" This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/) [ghstack-poisoned]

This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/) [ghstack-poisoned]

Pull Request resolved: #120806 This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. ghstack-source-id: 217307158 @exported-using-ghexport Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/)

This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. Differential Revision: [D54308741](https://our.internmc.facebook.com/intern/diff/D54308741/) [ghstack-poisoned]

Summary: This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. The next diff fails without this fix. ghstack-source-id: 217307158 exported-using-ghexport Test Plan: ``` buck2 run mode/opt caffe2/test:jit -- -r test_decorator ``` Differential Revision: D54308741

Summary: This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. Test Plan: ``` buck2 run mode/opt caffe2/test:jit -- -r test_decorator ``` Differential Revision: D54308741

Summary: This fixes a case left incomplete by #106229 The object is using __prepare_scriptable__ correctly inside of torch.jit.script() but the clousre that is obtained below is using the non-prepared version. This causes issues when the prepared and non-prepared versions are in different python modules. Test Plan: ``` buck2 run mode/opt caffe2/test:jit -- -r test_decorator ``` Differential Revision: D54308741 Re-exporting, as #120806 #121307 were not properly merged. Co-authored-by: Daniel Herrera <[email protected]> Pull Request resolved: #121553 Approved by: https://github.com/huydhn, https://github.com/seemethere

pytorch-bot bot added the release notes: jit release notes category label Jul 28, 2023

pytorchbot added the open source label Jul 28, 2023

pmeier requested review from ezyang and lezcano August 11, 2023 08:15

pmeier marked this pull request as ready for review August 11, 2023 08:15

lezcano approved these changes Aug 11, 2023

View reviewed changes

pmeier added 4 commits August 11, 2023 12:29

torch.jit.script escape hatch

d4ed92e

fix location

2a1e7f8

add test

56c9863

fix lint

1fcca1a

pmeier force-pushed the jit-escape branch from e69f69f to 1fcca1a Compare August 11, 2023 10:30

ezyang requested a review from suo August 11, 2023 15:23

ezyang approved these changes Aug 11, 2023

View reviewed changes

pytorchmergebot added the merging label Aug 11, 2023

pytorchmergebot added Merged and removed merging labels Aug 11, 2023

pytorchmergebot closed this in a926be3 Aug 11, 2023

pmeier deleted the jit-escape branch August 11, 2023 18:25

dherrera-meta mentioned this pull request Feb 28, 2024

[torch] Use __prepare_scriptable__ for closures #120806

Merged

dherrera-meta mentioned this pull request Mar 6, 2024

[torch] Use __prepare_scriptable__ for closures (#120806) #121307

Closed

izaitsevfb mentioned this pull request Mar 8, 2024

[torch] Use __prepare_scriptable__ for closures #121553

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

torch.jit.script escape hatch #106229

torch.jit.script escape hatch #106229

Uh oh!

pmeier commented Jul 28, 2023 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jul 28, 2023 •

edited

Loading

Uh oh!

lezcano left a comment

Uh oh!

ezyang left a comment

Uh oh!

pmeier commented Aug 11, 2023

Uh oh!

pytorchmergebot commented Aug 11, 2023

Uh oh!

Uh oh!

	if isinstance(obj, torch.nn.Module):
	obj = call_prepare_scriptable_func(obj)
	return torch.jit._recursive.create_script_module(
	obj, torch.jit._recursive.infer_methods_to_compile
	)

torch.jit.script escape hatch #106229

torch.jit.script escape hatch #106229

Uh oh!

Conversation

pmeier commented Jul 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/106229

❌ 1 New Failure

Uh oh!

lezcano left a comment

Choose a reason for hiding this comment

Uh oh!

ezyang left a comment

Choose a reason for hiding this comment

Uh oh!

pmeier commented Aug 11, 2023

Uh oh!

pytorchmergebot commented Aug 11, 2023

Merge started

Uh oh!

Uh oh!

pmeier commented Jul 28, 2023 •

edited

Loading

pytorch-bot bot commented Jul 28, 2023 •

edited

Loading