Fix wrong constant in fp8 (#38)

xadupre · web-flow · commit 4b6e04969d28 · 2023-08-01T18:04:33.000+02:00
* fix bug

* fix max value

* simplification
diff --git a/_unittests/ut_validation/test_f8.py b/_unittests/ut_validation/test_f8.py
@@ -924,7 +924,7 @@ def test_float8_e4m3fnuz_inf(self):
         x = numpy.float32(numpy.inf)
         to = float32_to_fe4m3(x, uz=True)
         back = fe4m3_to_float32(to, uz=True)
-        self.assertEqual(back, 224)
+        self.assertEqual(back, 240)
 
         x = numpy.float32(numpy.inf)
         to = float32_to_fe4m3(x, uz=True, saturate=False)
@@ -934,7 +934,7 @@ def test_float8_e4m3fnuz_inf(self):
         x = numpy.float32(-numpy.inf)
         to = float32_to_fe4m3(x, uz=True)
         back = fe4m3_to_float32(to, uz=True)
-        self.assertEqual(back, -224)
+        self.assertEqual(back, -240)
 
         x = numpy.float32(-numpy.inf)
         to = float32_to_fe4m3(x, uz=True, saturate=False)
diff --git a/onnx_array_api/validation/f8.py b/onnx_array_api/validation/f8.py
@@ -510,26 +510,22 @@ def float32_to_fe4m3(x, fn: bool = True, uz: bool = False, saturate: bool = True
             return 0x80
         if numpy.isinf(x):
             if saturate:
-                return ret | 126
+                return ret | 127
             return 0x80
         e = (b & 0x7F800000) >> 23  # exponent
         m = b & 0x007FFFFF  # mantissa
 
         if e != 0:
             if e < 116:
                 pass
-            elif e < 117:
-                # first positive number
-                if m > 0:
-                    ret |= 1
-                if (m >> 23) & 1:
-                    # rounding
-                    ret += 1
             elif e < 120:
                 # denormalized number
                 ex = e - 119
-                ret |= 1 << (2 + ex)
-                ret |= m >> (21 - ex)
+                if ex >= -2:
+                    ret |= 1 << (2 + ex)
+                    ret |= m >> (21 - ex)
+                elif m > 0:
+                    ret |= 1
                 mask = 1 << (20 - ex)
                 if m & mask and (
                     ret & 1
@@ -574,15 +570,14 @@ def float32_to_fe4m3(x, fn: bool = True, uz: bool = False, saturate: bool = True
         if e != 0:
             if e < 117:
                 pass
-            elif e < 118:
-                # first positive number
-                if m > 0:
-                    ret |= 1
             elif e < 121:
                 # denormalized number
                 ex = e - 120
-                ret |= 1 << (2 + ex)
-                ret |= m >> (21 - ex)
+                if ex >= -2:
+                    ret |= 1 << (2 + ex)
+                    ret |= m >> (21 - ex)
+                elif m > 0:
+                    ret |= 1
                 mask = 1 << (20 - ex)
                 if m & mask and (
                     ret & 1
@@ -642,18 +637,14 @@ def float32_to_fe5m2(x, fn: bool = False, uz: bool = False, saturate: bool = Tru
         if e != 0:
             if e < 109:
                 pass
-            elif e < 110:
-                # first positive number
-                if m > 0:
-                    ret |= 1
-                if (m >> 23) & 1:
-                    # rounding
-                    ret += 1
             elif e < 112:
-                # denormlized number
+                # denormalized number
                 ex = e - 111
-                ret |= 1 << (1 + ex)
-                ret |= m >> (22 - ex)
+                if ex >= -1:
+                    ret |= 1 << (1 + ex)
+                    ret |= m >> (22 - ex)
+                elif m > 0:
+                    ret |= 1
                 mask = 1 << (21 - ex)
                 if m & mask and (
                     ret & 1
@@ -696,18 +687,14 @@ def float32_to_fe5m2(x, fn: bool = False, uz: bool = False, saturate: bool = Tru
         if e != 0:
             if e < 110:
                 pass
-            elif e < 111:
-                # first positive number
-                if m > 0:
-                    ret |= 1
-                if (m >> 23) & 1:
-                    # rounding
-                    ret += 1
             elif e < 113:
-                # denormlized number
+                # denormalized number
                 ex = e - 112
-                ret |= 1 << (1 + ex)
-                ret |= m >> (22 - ex)
+                if ex >= -1:
+                    ret |= 1 << (1 + ex)
+                    ret |= m >> (22 - ex)
+                elif m > 0:
+                    ret |= 1
                 mask = 1 << (21 - ex)
                 if m & mask and (
                     ret & 1
