[autograd.Function] Add docs on the functorch interaction

[zou3519]

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

[pytorch-bot]

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

[soulitzer]

Nice!

[soulitzer]

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

[zou3519]

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.

[zou3519]

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

[soulitzer]

Will take another look once the docs render

[zou3519]

Will take another look once the docs render

Thanks!

[zou3519]

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

[zou3519]

@pytorchbot merge -f "unrelated failure"

[pytorchmergebot]

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