[quant][fix] at::print for per channel affine quantized tensors (#36280)

jerryzh168 · facebook-github-bot · commit f2bae8e86912 · 2020-04-09T16:27:00.000-07:00
Summary: Pull Request resolved: #36280 Test Plan: . Imported from OSS Differential Revision: D20948352 fbshipit-source-id: 92188806b9c129458ebb2cdc47599427e3b6e216
diff --git a/aten/src/ATen/core/Formatting.cpp b/aten/src/ATen/core/Formatting.cpp
@@ -278,8 +278,18 @@ std::ostream& print(std::ostream& stream, const Tensor & tensor_, int64_t linesi
     }
     if (tensor_.is_quantized()) {
       stream << ", qscheme: " << toString(tensor_.qscheme());
-      stream << ", scale: " << tensor_.q_scale();
-      stream << ", zero_point: " << tensor_.q_zero_point();
+      if (tensor_.qscheme() == c10::kPerTensorAffine) {
+        stream << ", scale: " << tensor_.q_scale();
+        stream << ", zero_point: " << tensor_.q_zero_point();
+      } else if (tensor_.qscheme() == c10::kPerChannelAffine) {
+        stream << ", scales: ";
+        Tensor scales = tensor_.q_per_channel_scales();
+        print(stream, scales, linesize);
+        stream << ", zero_points: ";
+        Tensor zero_points = tensor_.q_per_channel_zero_points();
+        print(stream, zero_points, linesize);
+        stream << ", axis: " << tensor_.q_per_channel_axis();
+      }
     }
     stream << " ]";
   }
