Updates to Python Module for Calculation of dX and Addition of Unit Tests (#42033)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/42033
In this diff, the Python `_LearnableFakeQuantize` module is updated where the gradient with respect to the input `x` is actually computed instead of passed through. Argument naming is also updated for better clarity; and unit tests on the `PerTensor` and `PerChannel` operators are added for asserting correctness.
Test Plan:
On a devvm, execute the command:
`buck test //caffe2/test:quantization -- learnable_py_module`
To include `cuda` tests as well, run:
`buck test mode/dev-nosan //caffe2/test:quantization -- learnable_py_module`
Reviewed By: jerryzh168
Differential Revision: D22735580
fbshipit-source-id: 66bea7e9f8cb6422936e653500f917aa597c86de
