Refactor get numerical jacobian (#54092)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/54092
This is the first of several refactors to get numerical jacobian:
This one just moves some logic around as to try to split the get_numerical_jacobian function into smaller more manageable functions:
- compute_gradient is now no longer nested, but we have to pass in the parameters instead
- iter_tensor extracts out the logic of iterating through different types of tensors (the code should be almost the exact same here except for instead of calling into the update jacobian function, we yield the arguments instead)
Test Plan: Imported from OSS
Reviewed By: H-Huang
Differential Revision: D27354268
Pulled By: soulitzer
fbshipit-source-id: 73288e3c889ae31bb8bf77a0e3acb3e9020e09a3
