fix float quant

Xilinx · Oct 8, 2024 · d3b520f · d3b520f
1 parent b3117dd
commit d3b520f
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/src/brevitas/core/quant/float.py b/src/brevitas/core/quant/float.py
@@ -68,12 +68,6 @@ def __init__(
 
  @brevitas.jit.script_method
  def quantize(self, x: torch.Tensor, scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
- if self.float_scaling_impl is not None:
- float_scaling_impl_value = self.float_scaling_impl(
- self.exponent_bit_width(), self.mantissa_bit_width(), self.exponent_bias())
- else:
- float_scaling_impl_value = None
- scale = self.scaling_impl(x, float_scaling_impl_value)
  x = self.input_view_impl(x)
  scaled_x = x / scale
  internal_scale = float_internal_scale(
@@ -87,7 +81,12 @@ def dequantize(self, y, scale):
 
  @brevitas.jit.script_method
  def forward(self, x):
- scale = self.scaling_impl(x)
+ if self.float_scaling_impl is not None:
+ float_scaling_impl_value = self.float_scaling_impl(
+ self.exponent_bit_width(), self.mantissa_bit_width(), self.exponent_bias())
+ else:
+ float_scaling_impl_value = None
+ scale = self.scaling_impl(x, float_scaling_impl_value)
  if self.observer_only:
  y = x
  saturating, inf_values, nan_values = self.float_clamp_impl.saturating, self.float_clamp_impl.inf_values, self.float_clamp_impl.nan_values

diff --git a/tests/brevitas/core/test_float_quant.py b/tests/brevitas/core/test_float_quant.py
@@ -143,7 +143,8 @@ def test_scaling_impls_called_once(inp, minifloat_format):
  scaling_impl=scaling_impl,
  float_scaling_impl=float_scaling_impl,
  float_clamp_impl=float_clamp)
- scale = float_quant.scaling_impl(inp)
+ float_scaling = float_scaling_impl(exponent_bit_width, mantissa_bit_width, exponent_bias)
+ scale = float_quant.scaling_impl(inp, float_scaling)
  _ = float_quant.quantize(inp, scale)
  # scaling implementations should be called exaclty once on the input
  float_scaling_impl.assert_called_once_with(