Add cachemask variant for fake_quantize_affine

[andrewor14]

Summary: In QAT, we often wish to filter out the gradients corresponding to values outside the expected quantization range, for example:

# Forward
q = _quantize_affine_no_dtype_cast(...)
dq = _dequantize_affine_no_dtype_check(...)
mask = torch.logical_and((q >= quant_min), (q <= quant_max))

# Backward
grad = grad * mask

The existing fake_quantize_affine returns the dequantized values only, so callers do not have access to this mask. This commit adds the variant to this op that returns both the dequantized values and the mask, similar to
fake_quantize_per_tensor_affine_cachemask in core.

Test Plan:
python test/quantization/test_quant_primitives.py -k test_fake_quantize_affine_cachemask

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/500

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit d70f92c with merge base aef7e09 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

we can dedup the comments by: https://github.com/pytorch/pytorch/blob/06ebf87a1eca6c345f5e3e39b63c2ef487695043/torch/ao/quantization/observer.py#L582

e.g.

:func:`~torchao.quantization.quant_primitives.fake_quantize_affine`

?

[jerryzh168]

please add a section for Args and link to fake_quantize_affine

[jerryzh168]

oh, I think we can move this to before Returns

[jerryzh168]

LG, had a comment on updating the docstring a bit