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

Xilinx · May 28, 2024 · 486a6df · 486a6df
1 parent b535615
commit 486a6df
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/brevitas/core/quant/float.py b/src/brevitas/core/quant/float.py
@@ -59,9 +59,14 @@ 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
+        if dtype is None:
+            dtype = torch.get_default_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