Repository: quic/aimet (language: python)

Issue:
Mark lazy quantize wrapper as unquantizable



Discussion / hints:

Fix patch:
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/quantsim.py b/TrainingExtensions/torch/src/python/aimet_torch/quantsim.py
index b2097fe3b7..984ed9b0ca 100644
--- a/TrainingExtensions/torch/src/python/aimet_torch/quantsim.py
+++ b/TrainingExtensions/torch/src/python/aimet_torch/quantsim.py
@@ -206,6 +206,7 @@ def get_original_module(self) -> torch.nn.Module:
     QcQuantizeRecurrent,
     ExportableQuantModule,
     torch.nn.Identity,
+    LazyQuantizeWrapper
 )
 
 


Test patch:
diff --git a/TrainingExtensions/torch/test/python/test_quantizer.py b/TrainingExtensions/torch/test/python/test_quantizer.py
index e5ef14f829..9a58ac5776 100644
--- a/TrainingExtensions/torch/test/python/test_quantizer.py
+++ b/TrainingExtensions/torch/test/python/test_quantizer.py
@@ -5158,6 +5158,37 @@ def test_fused_qdq_linear(input_dims):
 
     _assert_same_results_with_or_without_recompute(linear_wrapper, x)
 
+def test_multiple_references_to_same_module():
+    class TwoRelu(torch.nn.Module):
+        def __init__(self):
+            super(TwoRelu, self).__init__()
+            self.relu1 = torch.nn.ReLU()
+            self.relu2 = torch.nn.ReLU()
+
+        def forward(self, inp):
+            x = self.relu1(inp)
+            x = self.relu2(x)
+            return x
+
+    class OuterModel(torch.nn.Module):
+        def __init__(self):
+            super(OuterModel, self).__init__()
+            self.two_relu = TwoRelu()
+            self.seq = torch.nn.Sequential(self.two_relu)
+
+        def forward(self, inp):
+            x = self.seq(inp)
+            return x
+
+    model = OuterModel()
+    dummy_input = torch.randn(1, 2)
+
+    qsim = QuantizationSimModel(model, dummy_input)
+    qsim.compute_encodings(lambda m, _: m(dummy_input), None)
+    assert qsim.model.two_relu is qsim.model.seq[0]
+    assert isinstance(qsim.model.two_relu.relu1, StaticGridQuantWrapper)
+    assert isinstance(qsim.model.two_relu.relu1._module_to_wrap, torch.nn.ReLU)
+
 
 def _assert_same_results_with_or_without_recompute(wrapper: LearnedGridQuantWrapper, x):
     with no_recompute():