Fix mx.from_fp8 E4M3FN NaN decode with branchless carry

mlx/backend/cpu/unary_ops.h

@@ -154,19 +154,20 @@ struct ToFP8 {
 struct FromFP8 {
   template <int N>
   Simd<float, N> operator()(Simd<uint8_t, N> x) {
-    auto v = Simd<uint16_t, N>(x & 127) << 7;
+    auto v = Simd<uint16_t, N>(x & 127);
+    auto sign_bit = Simd<uint16_t, N>((x >> 7) & 1) << 15;
+    auto u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit;
     Simd<float, N> out;
     if constexpr (simd::max_size<float16_t> >= N) {
-      auto converted = *(Simd<float16_t, N>*)(&v);
+      auto converted = *(Simd<float16_t, N>*)(&u);
       out = converted * 256.0;
     } else {
       for (int i = 0; i < N; ++i) {
-        auto converted = *(float16_t*)(&v[i]);
+        auto converted = *(float16_t*)(&u[i]);
         out[i] = converted * 256.0;
       }
     }
-    auto sign = Simd<bool, N>(x & 128);
-    return select(sign, -out, out);
+    return out;
   }
   float operator()(uint8_t x) {
     return (*this)(Simd<uint8_t, 1>(x)).value;

mlx/backend/metal/kernels/fp8.h

@@ -30,11 +30,11 @@ struct fp8_e4m3 {
   }
 
   operator float16_t() thread {
-    uint16_t v = (bits & 127) << 7;
-    half converted = as_type<half>(v);
-    converted *= 256.0;
-    auto sign = bits & 128;
-    return (sign ? -converted : converted);
+    uint16_t v = bits & 127;
+    uint16_t sign_bit = ((uint16_t)((bits >> 7) & 1)) << 15;
+    uint16_t u = (v << 7) | (((v + 1) >> 7) << 14) | sign_bit;
+    half converted = as_type<half>(u);
+    return converted * 256.0;
   }
 
   operator bfloat16_t() thread {

python/tests/test_ops.py

@@ -4258,6 +4258,17 @@ def test_to_from_fp8(self):
         self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(vals)), vals))
         self.assertTrue(mx.array_equal(mx.from_fp8(mx.to_fp8(-vals)), -vals))
 
+        # 0x7f and 0xff are NaN in E4M3FN. Arithmetic does not keep the sign of
+        # a NaN, so only test for NaN. Use float32, numpy has no bfloat16.
+        nan_encodings = mx.array([0x7F, 0xFF], dtype=mx.uint8)
+        decoded_np = np.array(mx.from_fp8(nan_encodings, mx.float32))
+        self.assertTrue(np.isnan(decoded_np[0]).item())
+        self.assertTrue(np.isnan(decoded_np[1]).item())
+
+        # The carry must not fire one byte early: 0x7e/0xfe stay finite
+        finite = mx.from_fp8(mx.array([0x7E, 0xFE], dtype=mx.uint8), mx.float32)
+        self.assertTrue(mx.array_equal(finite, mx.array([448.0, -448.0])))
+
     def test_zeros_ones_empty_like_dtype(self):
         x = mx.array([1, 2, 3], dtype=mx.int32)