[autograd.Function] Add docs on the functorch interaction #91452

zou3519 · 2022-12-28T18:55:36Z

Stack from ghstack (oldest at bottom):

-> [autograd.Function] Add docs on the functorch interaction #91452

This PR:

Updates autograd.Function.forward docs to reflect how you either
define a forward with ctx or a separate forward and setup_context
Updates the "Extending Autograd" docs to suggest the usage of
autograd.Function with separate forward and setup_context. This should
be the default because there is a low barrier to go from this to
an autograd.Function that is fully supported by functorch transforms.
Adds a new "Extending torch.func with autograd.Function" doc that
explains how to use autograd.Function with torch.func. It also
explains how to use generate_vmap_rule and how to manually write a
vmap staticmethod.

While writing this, I noticed that the implementation of
setup_context staticmethod/generate_vmap_rule/vmap staticmethod are a
bit inconsistent with the other method/attributes on autograd.Function:

Asymmetry in implementation of autograd.Function and functorch interaction #91451
I'm happy to fix those if we think it is a problem, either in this PR
or a followup (this PR is getting long, I want some initial docs
out that I can point early adopters at, and fixing the problems in the
future isn't really BC-breaking).

Test Plan:

view docs preview

This PR: - Updates autograd.Function.forward docs to reflect how you either define a forward with ctx or a separate forward and setup_context - Updates the "Extending Autograd" docs to suggest the usage of autograd.Function with separate forward and setup_context. This should be the default because there is a low barrier to go from this to an autograd.Function that is fully supported by functorch transforms. - Adds a new "Extending torch.func with autograd.Function" doc that explains how to use autograd.Function with torch.func. It also explains how to use generate_vmap_rule and how to manually write a vmap staticmethod. While writing this, I noticed that the implementation of setup_context staticmethod/generate_vmap_rule/vmap staticmethod are a bit inconsistent with the other method/attributes on autograd.Function: - #91451 - I'm happy to fix those if we think it is a problem, either in this PR or a followup (this PR is getting long, I want some initial docs out that I can point early adopters at, and fixing the problems in the future isn't really BC-breaking). Test Plan: - view docs preview [ghstack-poisoned]

pytorch-bot · 2022-12-28T18:55:38Z

🔗 Helpful Links

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

📄 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.

❌ 2 Failures

As of commit 9334843:

NEW FAILURES - The following jobs have failed:

Labeler

FLAKY - The following jobs failed but were likely due to flakiness present on master:

linux-focal-rocm5.3-py3.8 / test (default, 2, 2, linux.rocm.gpu)

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

This PR: - Updates autograd.Function.forward docs to reflect how you either define a forward with ctx or a separate forward and setup_context - Updates the "Extending Autograd" docs to suggest the usage of autograd.Function with separate forward and setup_context. This should be the default because there is a low barrier to go from this to an autograd.Function that is fully supported by functorch transforms. - Adds a new "Extending torch.func with autograd.Function" doc that explains how to use autograd.Function with torch.func. It also explains how to use generate_vmap_rule and how to manually write a vmap staticmethod. While writing this, I noticed that the implementation of setup_context staticmethod/generate_vmap_rule/vmap staticmethod are a bit inconsistent with the other method/attributes on autograd.Function: - #91451 - I'm happy to fix those if we think it is a problem, either in this PR or a followup (this PR is getting long, I want some initial docs out that I can point early adopters at, and fixing the problems in the future isn't really BC-breaking). Test Plan: - view docs preview ghstack-source-id: 5b30efa Pull Request resolved: #91452

soulitzer

Nice!

docs/source/notes/extending.rst

torch/autograd/function.py

docs/source/notes/extending.rst

docs/source/notes/extending.func.rst

docs/source/notes/extending.rst

soulitzer · 2022-12-28T21:51:36Z

docs/source/notes/extending.func.rst

+Here's how to define the ``vmap`` staticmethod:
+
+- the signature is ``vmap(info, in_dims: Tuple[Optional[int]], *args)``, where
+  ``*args`` is the same as the args to :meth:`~Function.forward`.


should vmap take **kwargs for consistency since forward does (though they have to be passed positionally because apply doesn't take true kwargs)

I might take this as a follow-up item. I am not sure:

setup_context(ctx, inputs, output) takes in inputs that were directly passed to forward

If there were ANY unspecified values that are default kwargs, then we have a problem -- they're not passed to setup_context, so setup_context doesn't get their default value (if it needs to save them for backward)

So the conclusion is that autograd.Function with setup_context has some really awkward handling for default kwargs.

I'm inclined to say in the docs that autograd.Function with setup_context doesn't support **kwargs. A user should create a function wrapping their autograd.Function that actually does support **kwargs as the workaround.

I end up changing the documentation to suggest that default args and keywords args be passed via wrapper so it's clearer how to handle them, but not changing the suggested signature of the vmap staticmethod

docs/source/notes/extending.func.rst

soulitzer · 2022-12-29T01:53:47Z

Will take another look once the docs render

zou3519 · 2022-12-29T16:03:07Z

Will take another look once the docs render

Thanks!

This PR: - Updates autograd.Function.forward docs to reflect how you either define a forward with ctx or a separate forward and setup_context - Updates the "Extending Autograd" docs to suggest the usage of autograd.Function with separate forward and setup_context. This should be the default because there is a low barrier to go from this to an autograd.Function that is fully supported by functorch transforms. - Adds a new "Extending torch.func with autograd.Function" doc that explains how to use autograd.Function with torch.func. It also explains how to use generate_vmap_rule and how to manually write a vmap staticmethod. While writing this, I noticed that the implementation of setup_context staticmethod/generate_vmap_rule/vmap staticmethod are a bit inconsistent with the other method/attributes on autograd.Function: - #91451 - I'm happy to fix those if we think it is a problem, either in this PR or a followup (this PR is getting long, I want some initial docs out that I can point early adopters at, and fixing the problems in the future isn't really BC-breaking). Test Plan: - view docs preview [ghstack-poisoned]

This PR: - Updates autograd.Function.forward docs to reflect how you either define a forward with ctx or a separate forward and setup_context - Updates the "Extending Autograd" docs to suggest the usage of autograd.Function with separate forward and setup_context. This should be the default because there is a low barrier to go from this to an autograd.Function that is fully supported by functorch transforms. - Adds a new "Extending torch.func with autograd.Function" doc that explains how to use autograd.Function with torch.func. It also explains how to use generate_vmap_rule and how to manually write a vmap staticmethod. While writing this, I noticed that the implementation of setup_context staticmethod/generate_vmap_rule/vmap staticmethod are a bit inconsistent with the other method/attributes on autograd.Function: - #91451 - I'm happy to fix those if we think it is a problem, either in this PR or a followup (this PR is getting long, I want some initial docs out that I can point early adopters at, and fixing the problems in the future isn't really BC-breaking). Test Plan: - view docs preview ghstack-source-id: 1101af4 Pull Request resolved: #91452

zou3519 · 2022-12-29T19:57:29Z

Docs preview is up if folks wanted to take another pass: https://docs-preview.pytorch.org/91452/notes/extending.func.html

This PR: - Updates autograd.Function.forward docs to reflect how you either define a forward with ctx or a separate forward and setup_context - Updates the "Extending Autograd" docs to suggest the usage of autograd.Function with separate forward and setup_context. This should be the default because there is a low barrier to go from this to an autograd.Function that is fully supported by functorch transforms. - Adds a new "Extending torch.func with autograd.Function" doc that explains how to use autograd.Function with torch.func. It also explains how to use generate_vmap_rule and how to manually write a vmap staticmethod. While writing this, I noticed that the implementation of setup_context staticmethod/generate_vmap_rule/vmap staticmethod are a bit inconsistent with the other method/attributes on autograd.Function: - #91451 - I'm happy to fix those if we think it is a problem, either in this PR or a followup (this PR is getting long, I want some initial docs out that I can point early adopters at, and fixing the problems in the future isn't really BC-breaking). Test Plan: - view docs preview [ghstack-poisoned]

This PR: - Updates autograd.Function.forward docs to reflect how you either define a forward with ctx or a separate forward and setup_context - Updates the "Extending Autograd" docs to suggest the usage of autograd.Function with separate forward and setup_context. This should be the default because there is a low barrier to go from this to an autograd.Function that is fully supported by functorch transforms. - Adds a new "Extending torch.func with autograd.Function" doc that explains how to use autograd.Function with torch.func. It also explains how to use generate_vmap_rule and how to manually write a vmap staticmethod. While writing this, I noticed that the implementation of setup_context staticmethod/generate_vmap_rule/vmap staticmethod are a bit inconsistent with the other method/attributes on autograd.Function: - #91451 - I'm happy to fix those if we think it is a problem, either in this PR or a followup (this PR is getting long, I want some initial docs out that I can point early adopters at, and fixing the problems in the future isn't really BC-breaking). Test Plan: - view docs preview ghstack-source-id: be9a7e6 Pull Request resolved: #91452

zou3519 · 2023-01-04T00:26:33Z

@pytorchbot merge -f "unrelated failure"

pytorchmergebot · 2023-01-04T00:28:12Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes).

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

zou3519 requested review from albanD and soulitzer as code owners December 28, 2022 18:55

zou3519 mentioned this pull request Dec 28, 2022

[autograd.Function] enable the extended Function feature flag by default #91441

Closed

zou3519 requested review from bdhirsh and samdow December 28, 2022 18:56

zou3519 added release notes: autograd release notes category release notes: torch.func release notes category for torch.vmap or torch.func.* APIs labels Dec 28, 2022

soulitzer approved these changes Dec 28, 2022

View reviewed changes

zou3519 added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 3, 2023

pytorchmergebot added the Merged label Jan 4, 2023

pytorchmergebot closed this in 264f5ed Jan 4, 2023

facebook-github-bot deleted the gh/zou3519/598/head branch June 8, 2023 19:33

[autograd.Function] Add docs on the functorch interaction #91452

[autograd.Function] Add docs on the functorch interaction #91452

Uh oh!

Conversation

zou3519 commented Dec 28, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Dec 28, 2022 • 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/91452

❌ 2 Failures

Uh oh!

soulitzer left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

soulitzer Dec 28, 2022

Choose a reason for hiding this comment

Uh oh!

zou3519 Dec 29, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zou3519 Dec 29, 2022

Choose a reason for hiding this comment

Uh oh!

Uh oh!

soulitzer commented Dec 29, 2022

Uh oh!

zou3519 commented Dec 29, 2022

Uh oh!

zou3519 commented Dec 29, 2022

Uh oh!

zou3519 commented Jan 4, 2023

Uh oh!

pytorchmergebot commented Jan 4, 2023

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zou3519 commented Dec 28, 2022 •

edited

Loading

pytorch-bot bot commented Dec 28, 2022 •

edited

Loading

zou3519 Dec 29, 2022 •

edited

Loading