Fix (ptq): conflicts between gptq and equalization (#656)

src/brevitas/graph/gptq.py

@@ -227,6 +227,7 @@ def update_batch(self, module, input, current_layer):
         # Define batch size before re-organizing the input
         if hasattr(inp, 'names') and 'N' in inp.names:
             batch_dim = inp.names.index('N')
+            inp.rename_(None)
             inp = inp.transpose(0, batch_dim)
         batch_size = inp.shape[0]
 

src/brevitas/nn/equalized_layer.py

@@ -1,7 +1,11 @@
+from inspect import signature
+
 import torch
 
 from brevitas.nn.quant_mha import QuantMultiheadAttention
 
+INPUT_NAMES = ['input', 'inp', 'query', 'x']
+
 
 class EqualizedModule(torch.nn.Module):
 
@@ -11,9 +15,12 @@ def __init__(self, scale_module, layer) -> None:
         self.layer = layer
 
     def forward(self, *args, **kwargs):
-        kwargs.update(zip(self.layer.forward.__code__.co_varnames[1:], args))
+        # Convert args + kwargs + defaults into kwargs
+        bound_arguments = signature(self.layer.forward).bind(*args, **kwargs)
+        bound_arguments.apply_defaults()
+        kwargs = bound_arguments.arguments
 
-        possible_input_kwargs = ['input', 'inp', 'query']
+        possible_input_kwargs = INPUT_NAMES
         input_kwarg = [x for x in kwargs.keys() if x in possible_input_kwargs][0]
         x = kwargs[input_kwarg]
         out = x
@@ -31,5 +38,6 @@ def forward(self, *args, **kwargs):
         if isinstance(self.layer, (torch.nn.MultiheadAttention, QuantMultiheadAttention)):
             kwargs['key'] = out
             kwargs['value'] = out
-        out = self.layer(**kwargs)
+        # We convert everything to args so that hooks can work correctly
+        out = self.layer(*kwargs.values())
         return out