Fix (quant/float): restore fix for log(0)

Xilinx · Jun 10, 2024 · da2b89a · da2b89a
1 parent 02f5b6b
commit da2b89a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/src/brevitas/core/quant/float.py b/src/brevitas/core/quant/float.py
@@ -58,11 +58,6 @@ 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 quantize(self, x: torch.Tensor):
         scaling_impl_value = self.scaling_impl(x)

diff --git a/src/brevitas/utils/torch_utils.py b/src/brevitas/utils/torch_utils.py
@@ -95,8 +95,10 @@ def compute_channel_view_shape(tensor: torch.Tensor, channel_dim: int):
 def float_internal_scale(
         x: torch.Tensor, mantissa_bit_width: torch.Tensor,
         fp_internal_scale_min: torch.Tensor) -> torch.Tensor:
+    # Add small EPS to avoid log(0)
+    eps = torch.finfo(x.dtype).tiny
 
-    internal_scale = floor_ste(torch.log2(torch.abs(x))) - mantissa_bit_width
+    internal_scale = floor_ste(torch.log2(torch.abs(x) + eps)) - mantissa_bit_width
     internal_scale = torch.clamp_min(internal_scale, fp_internal_scale_min)
     internal_scale = torch.exp2(internal_scale)
     return internal_scale