Fix

Xilinx · Nov 19, 2024 · c399e9b · c399e9b
1 parent d9fff9a
commit c399e9b
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/brevitas/core/scaling/standalone.py b/src/brevitas/core/scaling/standalone.py
@@ -169,7 +169,6 @@ def __init__(
         self.restrict_clamp_threshold = _RestrictClampValue(
             restrict_value_impl=restrict_threshold_impl)
         self.restrict_threshold_pre = restrict_threshold_impl.restrict_init_module()
-        self.clamp_scaling = _ClampValue(scaling_min_val)
 
     @brevitas.jit.script_method
     def forward(self, placeholder: Tensor, threshold: Optional[Tensor] = None) -> Tensor:
@@ -178,9 +177,8 @@ def forward(self, placeholder: Tensor, threshold: Optional[Tensor] = None) -> Te
         # We first apply any restriction to scaling
         # For IntQuant, this is no-op, retrocompatible.
         threshold = self.restrict_clamp_threshold(self.restrict_threshold_pre(threshold))
-        # Clamping avoids eventual log(0) with restrict_val
-        value = self.clamp_scaling(self.value)
-        value = abs_binary_sign_grad(self.restrict_clamp_scaling(value))
+        # We can clamp after restrict val since the learned parameter is already in log-domain
+        value = abs_binary_sign_grad(self.restrict_clamp_scaling(self.value))
         return value / threshold
 
     def _load_from_state_dict(
@@ -426,6 +424,7 @@ def forward(self, stats_input: Tensor, threshold: Optional[Tensor] = None) -> Te
             threshold = self.restrict_threshold(self.restrict_threshold_pre(threshold))
             out = self.restrict_scaling(out)
             out = out / threshold
+            # We can clamp after restrict val since the learned parameter is already in log-domain
             out = abs_binary_sign_grad(self.clamp_scaling(out))
         return out