commit 9858f043508432a8c79691b93faf21d99d5cbf99
author Jerry Zhang <jerryzh168@gmail.com> Tue Sep 27 12:55:02 2022 -0700
committer PyTorch MergeBot <pytorchmergebot@users.noreply.github.com> Wed Sep 28 00:43:48 2022 +0000
tree 9b27ad04b736e80cb061d1cbe97183e450b170df
parent 7ff6a00a9a59ee53ca71dfa11697fb7822fd3c0c

[quant][docs] Add types for scale and zero_point tensor for torch.fake_quantize_per_channel_affine docs (#85733)

Summary:
Fixes: https://github.com/pytorch/pytorch/issues/85525

Test Plan:
visual inspection for the docs page

Reviewers:

Subscribers:

Tasks:

Tags:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85733
Approved by: https://github.com/z-a-f

torch/_torch_docs.py

diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
index 70c87c4..85de969 100644
--- a/torch/_torch_docs.py
+++ b/torch/_torch_docs.py
@@ -11927,14 +11927,14 @@
     )
 
 Args:
-    input (Tensor): the input value(s), in ``torch.float32``.
-    scale (double or Tensor): quantization scale
-    zero_point (int64 or Tensor): quantization zero_point
+    input (Tensor): the input value(s), ``torch.float32`` tensor
+    scale (double scalar or ``float32`` Tensor): quantization scale
+    zero_point (int64 scalar or ``int32`` Tensor): quantization zero_point
     quant_min (int64): lower bound of the quantized domain
     quant_max (int64): upper bound of the quantized domain
 
 Returns:
-    Tensor: A newly fake_quantized tensor
+    Tensor: A newly fake_quantized ``torch.float32`` tensor
 
 Example::
 
@@ -11966,15 +11966,15 @@
     )
 
 Args:
-    input (Tensor): the input value(s), in ``torch.float32``.
-    scale (Tensor): quantization scale, per channel
-    zero_point (Tensor): quantization zero_point, per channel
+    input (Tensor): the input value(s), in ``torch.float32``
+    scale (Tensor): quantization scale, per channel in ``torch.float32``
+    zero_point (Tensor): quantization zero_point, per channel in ``torch.int32`` or ``torch.half`` or ``torch.float32``
     axis (int32): channel axis
     quant_min (int64): lower bound of the quantized domain
     quant_max (int64): upper bound of the quantized domain
 
 Returns:
-    Tensor: A newly fake_quantized per channel tensor
+    Tensor: A newly fake_quantized per channel ``torch.float32`` tensor
 
 Example::
 
@@ -11988,7 +11988,7 @@
     >>> scales = (torch.randn(2) + 1) * 0.05
     >>> scales
     tensor([0.0475, 0.0486])
-    >>> zero_points = torch.zeros(2).to(torch.long)
+    >>> zero_points = torch.zeros(2).to(torch.int32)
     >>> zero_points
     tensor([0, 0])
     >>> torch.fake_quantize_per_channel_affine(x, scales, zero_points, 1, 0, 255)