Review

src/brevitas/core/stats/stats_op.py

@@ -12,6 +12,7 @@
 from brevitas import config
 from brevitas.core.utils import StatelessBuffer
 from brevitas.function.ops import max_int
+from brevitas.quant_tensor import _unpack_quant_tensor
 # Use custom implementation of kthvalue as work around to (b)float16 kernel limitations
 from brevitas.utils.torch_utils import kthvalue
 
@@ -478,8 +479,7 @@ def evaluate_loss(self, x, candidate):
  # Set to local_loss_mode before calling the proxy
  self.set_local_loss_mode(True)
  quant_value = self.proxy_forward(x)
- if isinstance(quant_value, tuple):
- quant_value = quant_value.value
+ quant_value = _unpack_quant_tensor(quant_value)
  loss = self.mse_loss_fn(x, quant_value)
  self.set_local_loss_mode(False)
  return loss

src/brevitas/nn/mixin/base.py

@@ -18,6 +18,7 @@
 from brevitas.inject import ExtendedInjector
 from brevitas.inject import Injector
 from brevitas.nn.utils import compute_channel_view_shape
+from brevitas.quant_tensor import _unpack_quant_tensor
 from brevitas.quant_tensor import QuantTensor
 
 from .utils import filter_kwargs
@@ -181,10 +182,7 @@ def pack_output(self, quant_output: QuantTensor):
  if self.return_quant_tensor:
  return quant_output
  else:
- if isinstance(quant_output, QuantTensor):
- return quant_output.value
- else:
- return quant_output
+ return _unpack_quant_tensor(quant_output)
 
 
 class QuantRecurrentLayerMixin(ExportMixin):
@@ -268,8 +266,6 @@ def maybe_quantize_input(self, inp):
  quant_input = inp
  if not self.quantize_output_only:
  quant_input = self.io_quant(quant_input)
- # elif not isinstance(inp, QuantTensor):
- # quant_input = QuantTensor(quant_input)
  return quant_input
 
  def maybe_quantize_state(self, inp, state, quant):

src/brevitas/nn/quant_layer.py

@@ -136,8 +136,7 @@ def forward(self, input: Union[Tensor, QuantTensor]):
  quant_input = self.input_quant(input)
  # shortcut execution through the export impl during export
  if self.export_mode:
- # quant_input_value = getattr(quant_input, 'value', quant_input)
- out = self.export_handler(quant_input)
+ out = self.export_handler(_unpack_quant_tensor(quant_input))
  self._set_global_is_quant_layer(False)
  return out
  out = self.act_quant(quant_input)
@@ -298,15 +297,14 @@ def merge_bn_in(self, bn):
  def forward_impl(self, inp: Union[Tensor, QuantTensor]) -> Union[Tensor, QuantTensor]:
  output_scale = None
  output_bit_width = None
- output_signed = None
  output_zero_point = None
+ output_signed = None
 
  inp = self.unpack_input(inp)
 
  # shortcut execution through the export impl during export
  if self.export_mode:
- inp_value = getattr(inp, 'value', inp)
- out = self.export_handler(inp_value)
+ out = self.export_handler(_unpack_quant_tensor(inp))
  self._set_global_is_quant_layer(False)
  return out
 
@@ -369,7 +367,7 @@ def forward_impl(self, inp: Union[Tensor, QuantTensor]) -> Union[Tensor, QuantTe
  elif self.return_quant_tensor and output_zero_point is None:
  output_zero_point = torch.zeros(1).type_as(output_tensor)
 
- if not self.return_quant_tensor or not compute_output_quant_tensor:
+ if not compute_output_quant_tensor:
  quant_output = output_tensor
  else:
  quant_output = QuantTensor(

src/brevitas/nn/quant_max_pool.py

@@ -7,6 +7,7 @@
 from torch.nn import MaxPool1d
 from torch.nn import MaxPool2d
 
+from brevitas.quant_tensor import _unpack_quant_tensor
 from brevitas.quant_tensor import QuantTensor
 
 from .mixin.base import QuantLayerMixin
@@ -81,8 +82,7 @@ def requires_export_handler(self):
  def forward(self, input: Union[Tensor, QuantTensor]):
  x = self.unpack_input(input)
  if self.export_mode:
- x_value = getattr(x, 'value', x)
- out = self.export_handler(x_value)
+ out = self.export_handler(_unpack_quant_tensor(x))
  self._set_global_is_quant_layer(False)
  return out
  x = x.set(value=super().forward(x.value))

src/brevitas/quant_tensor/__init__.py

@@ -20,10 +20,6 @@
 BFLOAT16_IS_VALID_ATOL = 0.5
 
 
-def _get_dequantize_tensor(input):
- return input.value if isinstance(input, QuantTensor) else input
-
-
 class QuantTensorBase(NamedTuple):
  value: Tensor
  scale: Optional[Tensor]

tests/brevitas/nn/test_linear.py

@@ -56,5 +56,7 @@ def test_forward_bias_int(self):
  torch.rand(size=(3, INPUT_FEATURES)),
  torch.tensor(1.0),
  torch.tensor(0.0),
- torch.tensor(3))
+ torch.tensor(3),
+ signed=True,
+ training=False)
  assert mod(x) is not None

tests/brevitas/nn/test_nn_quantizers.py

@@ -123,12 +123,12 @@ def test_quant_lstm_rnn_full(model_input, current_cases):
  assert isinstance(h, torch.Tensor)
 
  if c is not None:
- if not return_quant_tensor:
- assert isinstance(c, torch.Tensor)
- else:
+ if return_quant_tensor:
  assert isinstance(c, QuantTensor)
  assert c.scale is not None
  assert c.bit_width is not None
+ else:
+ assert isinstance(c, torch.Tensor)
 
 
 @pytest_cases.parametrize_with_cases('model_input', cases=[case_quant_lstm, case_quant_rnn])