Add backward compatible types to pt2e prepare (#2244)

metascroy · facebook-github-bot · commit c9b3d09daec9 · 2025-05-22T15:29:05.000-07:00
Summary: Pull Request resolved: pytorch/ao#2244 Differential Revision: D75248288
diff --git a/backends/cadence/aot/compiler.py b/backends/cadence/aot/compiler.py
@@ -123,7 +123,7 @@ def prepare_and_convert_pt2(
     assert isinstance(model_gm, torch.fx.GraphModule)
 
     # Prepare
-    prepared_model = prepare_pt2e(model_gm, quantizer)  # pyre-ignore[6]
+    prepared_model = prepare_pt2e(model_gm, quantizer)
 
     # Calibrate
     # If no calibration data is provided, use the inputs
diff --git a/examples/xnnpack/quantization/utils.py b/examples/xnnpack/quantization/utils.py
@@ -33,7 +33,7 @@ def quantize(
         is_dynamic=is_dynamic,
     )
     quantizer.set_global(operator_config)
-    m = prepare_pt2e(model, quantizer)  # pyre-ignore[6]
+    m = prepare_pt2e(model, quantizer)
     # calibration
     m(*example_inputs)
     m = convert_pt2e(m)
diff --git a/extension/llm/export/builder.py b/extension/llm/export/builder.py
@@ -373,7 +373,7 @@ def pt2e_quantize(self, quantizers: Optional[List[Quantizer]]) -> "LLMEdgeManage
                 ), "Please run export() first"
                 m = prepare_pt2e(
                     self.pre_autograd_graph_module,  # pyre-ignore[6]
-                    composed_quantizer,  # pyre-ignore[6]
+                    composed_quantizer,
                 )
                 logging.info(
                     f"Calibrating with tasks: {self.calibration_tasks}, limit: {self.calibration_limit}, calibration_data: {self.calibration_data}, tokenizer_path: {self.tokenizer_path}, seq_length: {self.calibration_seq_length}"

Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@ def quantize(`
`33`	`33`	`is_dynamic=is_dynamic,`
`34`	`34`	`)`
`35`	`35`	`quantizer.set_global(operator_config)`
`36`		`- m = prepare_pt2e(model, quantizer) # pyre-ignore[6]`
	`36`	`+ m = prepare_pt2e(model, quantizer)`
`37`	`37`	`# calibration`
`38`	`38`	`m(*example_inputs)`
`39`	`39`	`m = convert_pt2e(m)`
Original file line number	Diff line number	Diff line change
`@@ -373,7 +373,7 @@ def pt2e_quantize(self, quantizers: Optional[List[Quantizer]]) -> "LLMEdgeManage`
`373`	`373`	`), "Please run export() first"`
`374`	`374`	`m = prepare_pt2e(`
`375`	`375`	`self.pre_autograd_graph_module, # pyre-ignore[6]`
`376`		`- composed_quantizer, # pyre-ignore[6]`
	`376`	`+ composed_quantizer,`
`377`	`377`	`)`
`378`	`378`	`logging.info(`
`379`	`379`	`f"Calibrating with tasks: {self.calibration_tasks}, limit: {self.calibration_limit}, calibration_data: {self.calibration_data}, tokenizer_path: {self.tokenizer_path}, seq_length: {self.calibration_seq_length}"`