From 4f3920951eff8db6b71da99d9b92c3bb10679d64 Mon Sep 17 00:00:00 2001
From: Rong Rong <rongr@fb.com>
Date: Wed, 30 Sep 2020 09:15:13 -0700
Subject: [PATCH] type check for torch.quantization.quantize_jit (#45548)

Summary:
added type signal for more jit python functions

Pull Request resolved: https://github.com/pytorch/pytorch/pull/45548

Reviewed By: malfet

Differential Revision: D24010922

Pulled By: walterddr

fbshipit-source-id: 2fdd75482481adf2eddc01b915d7d5720fbb2b82
---
 mypy.ini                 |  3 ---
 torch/_C/__init__.pyi.in | 14 ++++++++++++++
 2 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/mypy.ini b/mypy.ini
index 24025938300..e8791f497c3 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -77,9 +77,6 @@ ignore_errors = True
 [mypy-torch.quantization.fake_quantize]
 ignore_errors = True
 
-[mypy-torch.quantization.quantize_jit]
-ignore_errors = True
-
 [mypy-torch.quantization._numeric_suite]
 ignore_errors = True
 
diff --git a/torch/_C/__init__.pyi.in b/torch/_C/__init__.pyi.in
index 4c2c65a33fe..b0479c01f58 100644
--- a/torch/_C/__init__.pyi.in
+++ b/torch/_C/__init__.pyi.in
@@ -179,6 +179,20 @@ def _jit_set_texpr_fuser_enabled(enable: _bool): ...
 def _jit_set_nvfuser_enabled(enable: _bool) -> _bool: ...
 def _jit_pass_canonicalize(graph: Graph): ...
 def _jit_pass_erase_shape_information(graph: Graph): ...
+def _jit_pass_fold_convbn(module: 'torch.jit.ScriptModule'): ...
+def _jit_pass_insert_observers(module: 'torch.jit.ScriptModule',
+                               method_name: str,
+                               qconfig_dict: Dict[str, Any],
+                               inplace: _bool,
+                               quant_type: _int): ...
+def _jit_pass_insert_quant_dequant(module: 'torch.jit.ScriptModule',
+                                   method_name: str,
+                                   inplace: _bool,
+                                   debug: _bool,
+                                   quant_type: _int): ...
+def _jit_pass_quant_finalize(module: 'torch.jit.ScriptModule',
+                             quant_type: _int,
+                             preserved_attrs: Sequence[str]): ...
 def _jit_set_profiling_executor(profiling_flag: _bool) -> _bool: ...
 def _jit_set_profiling_mode(profiling_flag: _bool) -> _bool: ...
 def _jit_try_infer_type(obj: Any) -> JitType: ...