Refactor functional api vectorized jacobian to use batched grad parameter (#65566)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/65566
This doesn't simplify vectorized jacobian computation, but is good to consolidate logic and helps us to test the logic
Test Plan: Imported from OSS
Reviewed By: albanD
Differential Revision: D31236257
Pulled By: soulitzer
fbshipit-source-id: 00ca0aa6519bed5f9ee2c7be4daa8872af5e92cd
