Add batched grad testing to OpInfo

[zou3519]

Stack from ghstack:

• Add batched grad testing to OpInfo #50818 Add batched grad testing to OpInfo

This PR does two things:

1. Add batched grad testing to OpInfo
2. Improve the error message from gradcheck if batched gradient
   computation fails to include suggestions for workarounds.

To add batched grad testing to OpInfo, this PR:

• adds new check_batched_grad=True and check_batched_gradgrad=True
  attributes to OpInfo. These are True by default because we expect most
  operators to support batched gradient computation.
• If check_batched_grad=True, then test_fn_grad invokes gradcheck
  with check_batched_grad=True.
• If check_batched_gradgrad=True, then test_fn_gradgradgrad invokes
  gradgradcheck with check_batched_grad=True.

The improved gradcheck error message looks like the following when an
exception is thrown while computing batched gradients:
https://gist.github.com/zou3519/5a0f46f908ba036259ca5e3752fd642f

Future

• Sometime in the not-near future, we will separate out "batched grad
  testing" from "gradcheck" for the purposes of OpInfo to make the
  testing more granular and also so that we can test that the vmap
  fallback doesn't get invoked (currently batched gradient testing only
  tests that the output values are correct).

Test Plan:

• run tests pytest test/test_ops.py -v -k "Gradients"

Differential Revision: D25997703

[albanD]

Looks good to me.

One question I would have about all these functions for which batched grad check fails is: should we make these hard errors for now to make sure they don't return silently wrong gradients?

[zou3519]

Looks good to me.

One question I would have about all these functions for which batched grad check fails is: should we make these hard errors for now to make sure they don't return silently wrong gradients?

My understanding is that they are hard errors -- are they actually soft errors in the code right now? (Is that what happens when gradcheck runs with raise_exception=False?)

[albanD]

My understanding is that they are hard errors

They are hard error when you run gradcheck because the gradients don't match (from my understanding, maybe I'm wrong here).
But if people try to use them, they will run without any error (and produce wrong gradients) <- I meant silent error here.

[facebook-github-bot]

@zou3519 merged this pull request in 1669151.