Fix (nn/bias): propagate runtime_shape from QuantScaleBias (#1385)

Author: Giuseppe5

src/brevitas/nn/quant_scale_bias.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 
 from typing import Optional
+from typing import Tuple
 from typing import Type
 from typing import Union
 
@@ -25,7 +26,8 @@
 
 class ScaleBias(Module):
 
-    def __init__(self, num_features: int, bias: bool, runtime_shape=(1, -1, 1, 1)):
+    def __init__(
+        self, num_features: int, bias: bool, runtime_shape: Tuple[int, ...] = (1, -1, 1, 1)):
         super(ScaleBias, self).__init__()
         self.num_features = num_features
         self.weight = Parameter(torch.ones(num_features))
@@ -49,9 +51,10 @@ def __init__(
             bias_quant: Optional[BiasQuantType] = None,
             input_quant: Optional[ActQuantType] = None,
             output_quant: Optional[ActQuantType] = None,
+            runtime_shape: Tuple[int, ...] = (1, -1, 1, 1),
             return_quant_tensor: bool = False,
             **kwargs) -> None:
-        ScaleBias.__init__(self, num_features, bias)
+        ScaleBias.__init__(self, num_features, bias, runtime_shape)
         QuantWBIOL.__init__(
             self,
             weight_quant=weight_quant,

tests/brevitas/nn/test_wbiol.py

@@ -12,7 +12,6 @@
 from brevitas.nn import QuantConvTranspose2d
 from brevitas.nn import QuantConvTranspose3d
 from brevitas.nn import QuantLinear
-from brevitas.nn import QuantScaleBias
 from brevitas.nn.quant_layer import QuantWeightBiasInputOutputLayer as QuantWBIOL
 from brevitas.proxy import ActQuantProxyFromInjector
 from brevitas.proxy import BiasQuantProxyFromInjector