[quant][graphmode] quantization support for view, transpose, contiguos (#34854)

jerryzh168 · facebook-github-bot · commit 537fdd77d5e9 · 2020-03-23T22:33:39.000-07:00
Summary: Pull Request resolved: #34854 Test Plan: python test/test_jit.py Imported from OSS Differential Revision: D20524456 fbshipit-source-id: e6e8fc3db6cccbd32c210d04f921274d81996fe2
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -2155,6 +2155,9 @@ def forward(self, x):
                 x = torch.min(x)
                 x = torch.mean(x)
                 x = x.reshape([-1])
+                x = x.view(-1)
+                x = x.transpose(1, 2)
+                x = x.contiguous()
                 x, y = torch.chunk(x, 2)
                 x = F.dropout(x)
                 x = self.dropout(x)
@@ -2178,6 +2181,9 @@ def forward(self, x):
                    .check("aten::min") \
                    .check("aten::mean") \
                    .check("aten::reshape") \
+                   .check("aten::view") \
+                   .check("aten::transpose") \
+                   .check("aten::contiguous") \
                    .check("aten::chunk") \
                    .check("prim::ListUnpack") \
                    .check("aten::dropout") \
@@ -2192,6 +2198,9 @@ def forward(self, x):
                    .check("aten::min") \
                    .check("aten::mean") \
                    .check("aten::reshape") \
+                   .check("aten::view") \
+                   .check("aten::transpose") \
+                   .check("aten::contiguous") \
                    .check("aten::chunk") \
                    .check("prim::ListUnpack") \
                    .check("aten::dropout") \
diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp
@@ -137,6 +137,9 @@ std::vector<size_t> getGeneralOpTensorInputIndexes(Node* n) {
       "dropout",
       "reshape",
       "chunk",
+      "view",
+      "transpose",
+      "contiguous",
       // TODO: sort returns a tuple of Tensors, we have
       // to extend the API to support that
       // "sort",
