Clean-up

Xilinx · Oct 26, 2024 · 475bd7f · 475bd7f
1 parent b4c9d34
commit 475bd7f
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/src/brevitas/quant_tensor/float_quant_tensor.py b/src/brevitas/quant_tensor/float_quant_tensor.py
@@ -103,9 +103,9 @@ def _pre_round_float_value(self):
             scale = self.scale.type(torch.float32)
         minifloat_value = value / scale
         fp_internal_scale = 1. - self.exponent_bias - self.mantissa_bit_width
-        eps = torch.finfo(self.scale.dtype).tiny
+        eps = torch.finfo(scale.dtype).tiny
         int_scale = float_internal_scale(
-            self.value / self.scale, self.mantissa_bit_width, fp_internal_scale, eps)
+            minifloat_value, self.mantissa_bit_width, fp_internal_scale, eps)
         minifloat_value = minifloat_value / int_scale
         return minifloat_value
 
@@ -138,12 +138,17 @@ def device(self):
     def minifloat(self, float_datatype=True):
         # TODO: Check if OCP and cast to proper data-type if matching
         assert float_datatype, "Minifloat quant returns only higher precision dtype"
-
         if self.is_valid:
+            value = self.value
+            scale = self.scale
+            if self.scale.dtype == torch.bfloat16:
+                value = self.value.type(torch.float32)
+                scale = self.scale.type(torch.float32)
+            minifloat_value = value / scale
             fp_internal_scale = 1. - self.exponent_bias - self.mantissa_bit_width
-            eps = torch.finfo(self.scale.dtype).tiny
+            eps = torch.finfo(scale.dtype).tiny
             int_scale = float_internal_scale(
-                self.value / self.scale, self.mantissa_bit_width, fp_internal_scale, eps)
+                minifloat_value, self.mantissa_bit_width, fp_internal_scale, eps)
             float_value = torch.round(self._pre_round_float_value) * int_scale
             return float_value.type(self.scale.dtype)
         else: