Add vectorized Jacobian and Hessian computation with forward AD (#67041)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/67041
Original PR here: https://github.com/pytorch/pytorch/pull/62246 (The old PR does more things, but now that's split across this stack)
This PR:
- Adds "jacfwd" and "hessian_fwdrev"
- Modifies existing tests to also test the `forward_ad=True` case
Test Plan: Imported from OSS
Reviewed By: gchanan, zou3519
Differential Revision: D32314424
Pulled By: soulitzer
fbshipit-source-id: 785b0e39162b93dc3b3cb9413233447152eddd53
