[quant][graphmode] quantization support for aten::chunk (#34806)

jerryzh168 · facebook-github-bot · commit 4a96911629bd · 2020-03-23T22:33:34.000-07:00
Summary: Pull Request resolved: #34806 Test Plan: python test/test_jit.py Imported from OSS Differential Revision: D20524454 fbshipit-source-id: 92ac9bc251581e963258cb90dc3de73f8508c822
diff --git a/test/test_jit.py b/test/test_jit.py
@@ -2155,6 +2155,7 @@ def forward(self, x):
                 x = torch.min(x)
                 x = torch.mean(x)
                 x = x.reshape([-1])
+                x, y = torch.chunk(x, 2)
                 x = F.dropout(x)
                 x = self.dropout(x)
                 # TODO: uncomment when sort is supported
@@ -2177,6 +2178,8 @@ def forward(self, x):
                    .check("aten::min") \
                    .check("aten::mean") \
                    .check("aten::reshape") \
+                   .check("aten::chunk") \
+                   .check("prim::ListUnpack") \
                    .check("aten::dropout") \
                    .check("aten::dropout") \
                    .run(m.graph)
@@ -2189,6 +2192,8 @@ def forward(self, x):
                    .check("aten::min") \
                    .check("aten::mean") \
                    .check("aten::reshape") \
+                   .check("aten::chunk") \
+                   .check("prim::ListUnpack") \
                    .check("aten::dropout") \
                    .check("aten::dropout") \
                    .check("dequantize") \
diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp
@@ -136,6 +136,7 @@ std::vector<size_t> getGeneralOpTensorInputIndexes(Node* n) {
       "upsample_bicubic2d",
       "dropout",
       "reshape",
+      "chunk",
       // TODO: sort returns a tuple of Tensors, we have
       // to extend the API to support that
       // "sort",
