From fcf732c61a71abc6e4368cd3b25edaa5fdfc82c1 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Tue, 29 Sep 2026 13:28:47 -0700 Subject: [PATCH] Arm backend: Add missing BasePipeline.set_quantization_calibration #23182 made the NSS and NFRU model tests call pipeline.set_quantization_calibration(), but never added the method, so the real-data TOSA INT tests fail with AttributeError. Add it to BasePipeline, configuring the quantize stage the same way the NSS test did directly before #23182. Authored with Claude Code. --- backends/arm/test/tester/test_pipeline.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/backends/arm/test/tester/test_pipeline.py b/backends/arm/test/tester/test_pipeline.py index 34a9ad62ea0..3ad4fd9cde5 100644 --- a/backends/arm/test/tester/test_pipeline.py +++ b/backends/arm/test/tester/test_pipeline.py @@ -11,6 +11,7 @@ Callable, Dict, Generic, + Iterable, List, Optional, Sequence, @@ -246,6 +247,20 @@ def quantizer(self) -> TOSAQuantizer: f"First argument of quantize stage was {type(quantize_stage).__name__}, not Quantize as expected." ) + def set_quantization_calibration( + self, + calibration_samples: Iterable[Any], + dynamic_shapes: Optional[Tuple[Any, ...]] = None, + ): + """Calibrates the quantize stage with the given samples instead of the + test data. + """ + quantize_stage = self._stages[self.find_pos("quantize")].args[0] + quantize_stage.calibration_samples = calibration_samples + if dynamic_shapes is not None: + quantize_stage.dynamic_shapes = dynamic_shapes + return self + def pop_stage(self, identifier: int | str): """Removes and returns the stage at postion pos.""" if isinstance(identifier, int):