[reland] Add batched grad testing to gradcheck, turn it on in test_autograd

[zou3519]

Stack from ghstack:

• [reland] Add batched grad testing to gradcheck, turn it on in test_autograd #50592 [reland] Add batched grad testing to gradcheck, turn it on in test_autograd

This adds a check_batched_grad=False option to gradcheck and gradgradcheck.
It defaults to False because gradcheck is a public API and I don't want
to break any existing non-pytorch users of gradcheck.
This:

• runs grad twice with two grad outputs, a & b
• runs a vmapped grad with torch.stack([a, b])
• compares the results of the above against each other.

Furthermore:

• check_batched_grad=True is set to be the default for
  gradcheck/gradgradcheck inside of test_autograd.py. This is done by
  reassigning to the gradcheck object inside test_autograd
• I manually added check_batched_grad=False to gradcheck instances
  that don't support batched grad.
• I added a denylist for operations that don't support batched grad.

Question:

• Should we have a testing only gradcheck (e.g.,
  torch.testing.gradcheck) that has different defaults from our public
  API, torch.autograd.gradcheck?

Future:

• The future plan for this is to repeat the above for test_nn.py (the
  autogenerated test will require a denylist)
• Finally, we can repeat the above for all pytorch test files that use
  gradcheck.

Test Plan:

• run tests

Differential Revision: D25925942

[zou3519]

Note to reviewers: this is a re-land of #49120, which broke a slow test on master. This reland turns off batch grad testing for that specific test

[albanD]

👍

[codecov]

Codecov Report

Merging #50592 (55d7811) into gh/zou3519/339/base (7f3a407) will increase coverage by 0.00%.
The diff coverage is 92.30%.

@@                 Coverage Diff                  @@
##           gh/zou3519/339/base   #50592   +/-   ##
====================================================
  Coverage                80.66%   80.66%           
====================================================
  Files                     1913     1913           
  Lines                   208058   208083   +25     
====================================================
+ Hits                    167833   167854   +21     
- Misses                   40225    40229    +4     

[facebook-github-bot]

@zou3519 merged this pull request in f7a8bfd.