Update on "[quant][eagermode] Move custom_module registration to prep…

…are/convert_custom_config_dict"

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: [D24290811](https://our.internmc.facebook.com/intern/diff/D24290811)

[ghstack-poisoned]

tools/jit/gen_unboxing_wrappers.py

@@ -21,6 +21,7 @@
 import argparse
 import re
 from itertools import groupby
+from functools import reduce
 from ..autograd.gen_autograd import load_aten_declarations
 from ..autograd.gen_autograd import RETURNS_VIEWS_OF_INPUT
 from ..autograd.utils import CodeTemplate, write, is_out_variant, op_name_without_overload
@@ -496,6 +497,15 @@ def expand_options(decl, i, arg):
         }
         write(out, 'generated_unboxing_wrappers_%d.cpp' % i, GENERATED_UNBOXING_WRAPPERS_CPP, env)
 
+    all_shards = reduce(
+        lambda lhs, rhs: lhs + rhs,
+        shards,
+    )
+    env = {
+        'constructors': all_shards,
+    }
+    write(out, 'generated_unboxing_wrappers_everything.cpp', GENERATED_UNBOXING_WRAPPERS_CPP, env)
+
 
 default_map = {'{}': 'None', 'nullptr': 'None', 'c10::nullopt': 'None'}
 

torch/quantization/fx/quantization_patterns.py

@@ -573,7 +573,7 @@ class StandaloneModuleQuantizeHandler(QuantizeHandler):
     """ Converts an observed standalone module to quantized standalone module
     by calling convert_fx on the observed standalone module.
     """
-    def convert(self, quantizer, node, load_arg, debug=False):
+    def convert(self, quantizer, node, load_arg, debug=False, convert_custom_config_dict=None):
         assert node.op == 'call_module'
         qconfig = quantizer.qconfig_map[node.name]
         convert = torch.quantization.quantize_fx._convert_standalone_module_fx