Fix (GPxQ): unwrap QuantTensor when dealing with QuantLinear

src/brevitas/graph/gpfq.py

@@ -168,6 +168,9 @@ def update_batch(self, module, input, current_layer):
         if isinstance(self.layer, qnn.QuantLinear):
             if len(inp.shape) > 2:
                 inp = inp.reshape((-1, sum(inp.shape[2:])))
+            # Unwrap tensor value if quantized input
+            if isinstance(inp, QuantTensor):
+                inp = inp.value
             # For QuantLinear layer, groups will be 1
             inp_processed = inp.unsqueeze(0)
 

src/brevitas/graph/gptq.py

@@ -21,6 +21,7 @@
 from brevitas.graph.gpxq import StopFwdException
 from brevitas.graph.gpxq import SUPPORTED_CONV_OP
 import brevitas.nn as qnn
+from brevitas.quant_tensor import QuantTensor
 
 
 class gptq_mode(gpxq_mode):
@@ -150,6 +151,9 @@ def update_batch(self, module, input, current_layer):
         if isinstance(self.layer, qnn.QuantLinear):
             if len(inp.shape) > 2:
                 inp = inp.reshape((-1, sum(inp.shape[2:])))
+            # Unwrap tensor value if QuantTensor
+            if isinstance(inp, QuantTensor):
+                inp = inp.value
             inp = inp.t()
             # For QuantLinear layer, groups will be 1
             inp_processed = inp.unsqueeze(0)

tests/brevitas/graph/equalization_fixtures.py

@@ -461,11 +461,31 @@ def forward(self, x):
     return QuantConvTransposeModel
 
 
+@pytest_cases.fixture()
+def quant_linear_model():
+
+    class QuantLinearModel(nn.Module):
+
+        def __init__(self) -> None:
+            super().__init__()
+            self.linear_0 = qnn.QuantLinear(
+                3, 16, True, input_quant=Int8ActPerTensorFloat, return_quant_tensor=True)
+            self.linear_1 = qnn.QuantLinear(16, 10, True)
+
+        def forward(self, x):
+            x = self.linear_0(x)
+            x = self.linear_1(x)
+            return x
+
+    return QuantLinearModel
+
+
 list_of_quant_fixtures = [
     'quant_conv_with_input_quant_model',
     'quant_convdepthconv_model',
     'quant_residual_model',
-    'quant_convtranspose_model']
+    'quant_convtranspose_model',
+    'quant_linear_model']
 
 toy_quant_model = fixture_union(
     'toy_quant_model', list_of_quant_fixtures, ids=list_of_quant_fixtures)

tests/brevitas/graph/test_gpxq.py

@@ -111,7 +111,7 @@ def test_toymodels(
 
     model_class = toy_quant_model
     model = model_class()
-    if 'mha' in test_id:
+    if 'mha' in test_id or 'linear' in test_id:
         inp = torch.randn(32, *IN_SIZE_LINEAR[1:])
     else:
         inp = torch.randn(32, *IN_SIZE_CONV_SMALL[1:])