Fix (core/quant/float): use eps to avoid log(0)

Xilinx · May 23, 2024 · 8a9196b · 8a9196b
1 parent eeffd2b
commit 8a9196b
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/brevitas/core/quant/float.py b/src/brevitas/core/quant/float.py
@@ -59,9 +59,12 @@ def __init__(
         self.scaling_impl = scaling_impl
         self.float_clamp_impl = float_clamp_impl
 
+        # To avoid log(0), we add small a small value based on the used dtype
+        self.eps = torch.finfo(dtype).tiny
+
     @brevitas.jit.script_method
     def internal_scale(self, x):
-        internal_scale = floor_ste(torch.log2(torch.abs(x))) - self.mantissa_bit_width()
+        internal_scale = floor_ste(torch.log2(torch.abs(x) + self.eps)) - self.mantissa_bit_width()
         internal_scale = torch.clamp_min(internal_scale, self.fp_internal_scale_min())
         internal_scale = torch.exp2(internal_scale)
         return internal_scale