Fix (core/float): add default for float_scaling_impl (#972)

Xilinx · Jun 19, 2024 · df1a137 · df1a137
1 parent 5b9b1f8
commit df1a137
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/src/brevitas/core/quant/float.py b/src/brevitas/core/quant/float.py
@@ -65,10 +65,13 @@ def __init__(
 
  @brevitas.jit.script_method
  def quantize(self, x: torch.Tensor):
- scaling_impl_value = self.scaling_impl(x)
- float_scaling_impl_value = self.float_scaling_impl(
- self.exponent_bit_width(), self.mantissa_bit_width(), self.exponent_bias())
- scale = scaling_impl_value / float_scaling_impl_value
+ 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())
+ scale = scale / float_scaling_impl_value
+
  scaled_x = x / scale
  internal_scale = float_internal_scale(
  scaled_x, self.mantissa_bit_width(), self.fp_internal_scale_min(), self.eps)